import time
import uuid
import neat
import uuid
import numpy as np
import vrep.vrep as vrep
import threading
import pickle
from datetime import datetime, timedelta
from vision.tracker import get_marker_object
from robot.vrep_robot import VrepRobot
from utility.path_tracking import follow_path, transform_pos_angle, create_grid
from utility.util_functions import scale, euclidean_distance, \
    f_wheel_center, f_straight_movements, \
    f_obstacle_dist, scale, scale_thymio_sensors, \
    normalize_0_1, f_t_obstacle_avoidance, thymio_position, \
    flatten_dict, calc_behavioral_features, save_debug_data
try:
    from robot.evolved_robot import EvolvedRobot
except ImportError as error:
    print(error.__class__.__name__ + ": " + 'DBus works only on linux!')
from multiprocessing import current_process
from vrep.control_env import get_object_handle, get_pose, set_pose
import schedule

def eval_genomes_hardware(individual, settings, genomes, config, generation):
    """Evaluation function to evaluate NEAT genomes on Thymio robot"""

    for _, genome in genomes:
        _ = eval_genome_hardware(
            individual,
            settings,
            genome,
            model=None,
            config=config,
            generation=generation
        )

def eval_genomes_simulation(individual, settings, genomes, config, generation):
    """Evaluation function to evaluate NEAT genomes in VREP simulator"""

    for _, genome in genomes:
        _ = eval_genome_simulation(
            individual,
            settings,
            None,
            config,
            generation,
            genome
        )

def post_eval_genome(individual, settings, genome, model=None, config=None, generation=None):
    """Post evaluation of controllers using NEAT.
       Only used for testing controllers evolved using NEAT.

       :individual: `VrepRobot` or `EvolvedRobot` (Thymio)
       :genome: controller
       :config: NN configuration
    """

    print('Postevaluation of {0} started!'.format(type(individual).__name__))

    # neural network initialization
    network = init_network(genome, config, model)

    if type(individual).__name__ == 'VrepRobot':
        individual.v_chromosome = genome
        individual.id = genome.key

        # Enable the synchronous mode
        vrep.simxSynchronous(settings.client_id, True)
        if (vrep.simxStartSimulation(settings.client_id, vrep.simx_opmode_oneshot) == -1):
            return

        # collistion detection initialization
        collision_handle, collision = init_collision(individual.client_id)

        dt = 0.05
        runtime = 0
        steps = 0

        while not collision and settings.run_time > runtime:
            # The first simulation step waits for a trigger before being executed
            vrep.simxSynchronousTrigger(settings.client_id)
            _, collision = vrep.simxReadCollision(
                settings.client_id, collision_handle, vrep.simx_opmode_buffer)

            individual.v_neuro_loop()

            # input data to the neural network
            if type(network).__name__ == 'FeedForwardNetwork':
                net_output = network.activate(
                    individual.v_norm_sensor_activation)

            if type(network).__name__ == 'Sequential':
                net_output = network.predict(
                    (individual.v_norm_sensor_activation).reshape((1, 7)))[0]
            # [-2, 2] wheel speed thymio
            scaled_output = np.array(
                [scale(xi, -2.0, 2.0) for xi in net_output])
            # set motor wheel speeds
            individual.v_set_motors(*list(scaled_output))
            # After this call, the first simulation step is finished
            vrep.simxGetPingTime(settings.client_id)
            runtime += dt
            steps += 1

        # Before closing the connection to V-REP, make sure that the last command sent out had time to arrive.
        vrep.simxGetPingTime(settings.client_id)

        if (vrep.simxStopSimulation(settings.client_id, settings.op_mode) == -1):
            return
        return individual

    elif type(individual).__name__ == 'EvolvedRobot':
        init_position = np.array([0.19, 0.22])

        # Build Scene and get the obstacles grid
        obstacle_grid = build_scene(
            settings.config_scene, settings, individual.client_id)

        individual.chromosome = genome
        individual.id = genome.key
        t_xy, t_angle = thymio_position()

        if (vrep.simxStartSimulation(settings.client_id, vrep.simx_opmode_oneshot) == -1):
            print('Failed to start the simulation\n')
            return

        # update position and orientation of the robot in vrep
        position, orientation = transform_pos_angle(
            t_xy, t_angle)
        individual.v_set_pos_angle(position, orientation)

        # collistion detection initialization
        collision_handle, collision = init_collision(individual.client_id)

        now = datetime.now()

        while datetime.now() - now < timedelta(seconds=settings.run_time):

            t_xy, t_angle = thymio_position()
            # update position and orientation of the robot in vrep
            position, orientation = transform_pos_angle(
                t_xy, t_angle)
            individual.v_set_pos_angle(position, orientation)

            _, collision = vrep.simxReadCollision(
                settings.client_id, collision_handle, vrep.simx_opmode_buffer)
            # read proximity sensors data
            individual.t_read_prox()

            # input data to the neural network
            if type(network).__name__ == 'FeedForwardNetwork':
                net_output = network.activate(individual.n_t_sensor_activation)

            if type(network).__name__ == 'Sequential':
                net_output = network.predict(
                    (individual.n_t_sensor_activation).reshape((1, 7)))[0]

            # normalize motor wheel wheel_speeds [0.0, 2.0] - robot
            scaled_output = np.array([scale(xi, -200, 200)
                                      for xi in net_output])
            # set thymio wheel speeds
            individual.t_set_motors(*list(scaled_output))

        individual.t_stop()

        follow_path(
            individual,
            init_position,
            get_marker_object,
            vrep,
            settings.client_id,
            grid=obstacle_grid,
            log_time=settings.logtime_data
        )

        if (vrep.simxStopSimulation(settings.client_id, settings.op_mode) == -1):
            print('Failed to stop the simulation')
            print('Program ended')
            return
        return individual
    else:
        return None

def eval_genome_hardware(individual, settings, genome, model=None, config=None, generation=None):
    """Evaluation function for a single genome encoded with NEAT or DEAP."""

    robot_m = get_marker_object(7)
    while robot_m.realxy() is None:
        # obtain goal marker postion
        robot_m = get_marker_object(7)
    init_position = np.array([0.19, 0.22])

    # Build Scene and get the obstacles grid
    obstacle_grid = build_scene(
        settings.config_scene, settings, individual.client_id)

    # individual reset
    individual.n_t_sensor_activation = np.array([])
    individual.chromosome = genome
    individual.id = genome.key

    # simulation specific props
    thymio_position = [np.array([0.19, 0.22])]
    schedule.every(2).seconds.do(thymio_get_position_every_2s, thymio_position)

    collision = False
    scaled_output = np.array([])
    fitness_agg = np.array([])

    dt = 0.05
    runtime = 0
    steps = 0
    # Behavioral Features #1 and #2
    wheel_speeds = []
    sensor_activations = []

    # neural network initialization
    network = init_network(genome, config, model)

    # get robot marker
    robot_m = get_marker_object(7)
    if robot_m.realxy() is not None:
        # update current position of the robot
        robot_current_position = robot_m.realxy()[:2]

    # update position and orientation of the robot in vrep
    pos, orientation = transform_pos_angle(
        robot_current_position, robot_m.orientation())
    individual.v_set_pos_angle(pos, orientation)

    if (vrep.simxStartSimulation(individual.client_id, vrep.simx_opmode_oneshot) == -1):
        print('Failed to start the simulation\n')
        return

    # collistion detection initialization
    collision_handle, collision = init_collision(individual.client_id)

    # areas detection initlaization
    areas_name = ('area0', 'area1', 'area2')
    areas_handle = [(area,) + vrep.simxGetCollisionHandle(
        individual.client_id, area, vrep.simx_opmode_blocking) for area in areas_name]

    _ = [(handle[0],) + vrep.simxReadCollision(
        individual.client_id, handle[2], vrep.simx_opmode_streaming) for handle in areas_handle]

    areas_counter = dict([(area, dict(count=0, percentage=0.0, total=0))
                          for area in areas_name])

    now = datetime.now()

    # while datetime.now() - now < timedelta(seconds=genome.sim_time):
    while not collision and datetime.now() - now < timedelta(seconds=genome.sim_time):
        schedule.run_pending()
        # get robot marker
        robot_m = get_marker_object(7)
        if robot_m.realxy() is not None:
            # update current position of the robot
            robot_current_position = robot_m.realxy()[:2]

        # update position and orientation of the robot in vrep
        pos, orientation = transform_pos_angle(
            robot_current_position, robot_m.orientation())
        individual.v_set_pos_angle(pos, orientation)

        _, collision = vrep.simxReadCollision(
            individual.client_id, collision_handle, vrep.simx_opmode_buffer)

        # Behavioral Feature #3
        areas = [(handle[0],) + vrep.simxReadCollision(
            individual.client_id, handle[2], vrep.simx_opmode_streaming) for handle in areas_handle]

        for area, _, detected in areas:
            if detected:
                areas_counter.get(area).update(
                    count=areas_counter.get(area)['count']+1)

        # read proximity sensors data
        individual.t_read_prox()

        # input data to the neural network
        if type(network).__name__ == 'FeedForwardNetwork':
            net_output = network.activate(individual.n_t_sensor_activation)

        if type(network).__name__ == 'Sequential':
            net_output = network.predict(
                (individual.n_t_sensor_activation).reshape((1, 7)))[0]

        # normalize motor wheel wheel_speeds [0.0, 2.0] - robot
        scaled_output = np.array([scale(xi, -200, 200)
                                  for xi in net_output])

        # Collect behavioral feature data
        wheel_speeds.append(net_output)
        sensor_activations.append(
            list(map(lambda x: 1 if x > 0.0 else 0, individual.n_t_sensor_activation)))

        # set thymio wheel speeds
        individual.t_set_motors(*list(scaled_output))

        runtime += dt
        steps += 1

        # every 10 seconds the robot is in the same position given a threshold stop the simulation
        #  if round(runtime, 2) % 10.0 == 0.0:
        #      print(euclidean_distance(thymio_position[0], thymio_position[-1]))
        #      if (euclidean_distance(thymio_position[0], thymio_position[-1])) < .09:
        #          collision = True
        #
        #  fitness_t at time stamp
        (
            fitness_t,
            wheel_center,
            straight_movements,
            obstacles_distance
        ) = f_t_obstacle_avoidance(
            scaled_output, individual.n_t_sensor_activation, 'thymio')

        fitness_agg = np.append(fitness_agg, fitness_t)

        # dump individual data
        if settings.debug:
            save_debug_data(
                settings.path,
                individual.id,
                individual.t_sensor_activation,
                individual.n_t_sensor_activation,
                net_output,
                scaled_output,
                wheel_center,
                straight_movements,
                obstacles_distance,
                fitness_t,
                'THYMIO',
                robot_current_position
            )

    individual.t_stop()
    # calculate the fitnesss
    fitness = np.sum(fitness_agg)/settings.run_time

    schedule.clear()

    print('Transfered to thymio genome_id: {} fitness: {:.4f} runtime: {:.2f} s steps: {}'.format(
        individual.id, fitness, runtime, steps))

    if type(genome).__name__ == 'Individual':
        generation = genome.gen

    behavioral_features = calc_behavioral_features(
        areas_counter,
        wheel_speeds,
        sensor_activations,
        settings.path,
        genome.key,
        generation,
        'THYMIO'
    )

    if settings.debug:
        print('behavioral_features: {0}\n pos_sample: {1}\n'.format(
            behavioral_features, thymio_position))

    if type(genome).__name__ == 'Individual':
        genome.features = behavioral_features
        genome.task_fitness = fitness
        genome.position = thymio_position
        genome.evaluation = 'THYMIO'
        genome.weights = network.get_weights()

    if type(genome).__name__ == 'DefaultGenome':
        genome.features = behavioral_features
        genome.fitness = fitness
        genome.position = thymio_position

    follow_path(
        individual,
        init_position,
        get_marker_object,
        vrep,
        individual.client_id,
        grid=obstacle_grid,
        log_time=settings.logtime_data
    )

    if (vrep.simxStopSimulation(individual.client_id, settings.op_mode) == -1):
        print('Failed to stop the simulation')
        print('Program ended')
        return

    time.sleep(1)

    return fitness

def eval_genome_simulation(individual, settings, model, config, generation, genome):
    """Evaluation function for multiobjective optimization NSGA-II.
       :individual: robotic controller (VREP)
       :settings: simulation specific settings
       :model: Keras model Feedforward NN
       :genome: weights of the NN encoded that are being optimized

       :return: (fitness, transferability)

       fitness - task dependent fitness value V * (1 - sqr(delta v)) * (1 - max(S_activation))
       transferability - measure the distance betweeen simulation and real behavior
    """

    # reset the individual
    individual.v_reset_init()
    individual.chromosome = genome
    individual.id = genome.key

    # Behavioral Features #1 and #2
    wheel_speeds = []
    sensor_activations = []

    position = []

    # evaluation specific props
    collision = False
    scaled_output = np.array([], ndmin=2)
    fitness_agg = np.array([], ndmin=2)

    # setting time step to 50 ms (miliseconds)
    dt = 0.05
    runtime = 0
    steps = 0

    network = init_network(genome, config, model)

    # Enable the synchronous mode
    vrep.simxSynchronous(individual.client_id, True)

    # start the simulation
    if (vrep.simxStartSimulation(individual.client_id, vrep.simx_opmode_oneshot) == -1):
        return

    # collistion detection initialization
    collision_handle, collision = init_collision(individual.client_id)

    # areas detection initlaization
    areas_name = ('area0', 'area1', 'area2')
    areas_handle = [(area,) + vrep.simxGetCollisionHandle(
        individual.client_id, area, vrep.simx_opmode_blocking) for area in areas_name]

    _ = [(handle[0],) + vrep.simxReadCollision(
        individual.client_id, handle[2], vrep.simx_opmode_streaming) for handle in areas_handle]

    areas_counter = dict([(area, dict(count=0, percentage=0.0, total=0))
                          for area in areas_name])

    # while settings.run_time > runtime:
    while not collision and settings.run_time > runtime:
        # The first simulation step waits for a trigger before being executed
        vrep.simxSynchronousTrigger(individual.client_id)

        # read the collision
        _, collision = vrep.simxReadCollision(
            individual.client_id, collision_handle, vrep.simx_opmode_buffer)

        # Behavioral Feature #3
        areas = [(handle[0],) + vrep.simxReadCollision(
            individual.client_id, handle[2], vrep.simx_opmode_streaming) for handle in areas_handle]

        for area, _, detected in areas:
            if detected:
                areas_counter.get(area).update(
                    count=areas_counter.get(area)['count']+1)

        # read proximity sensors data
        individual.v_neuro_loop()

        # input data to the neural network
        if type(network).__name__ == 'FeedForwardNetwork':
            net_output = network.activate(individual.v_norm_sensor_activation)

        if type(network).__name__ == 'Sequential':
            net_output = network.predict(
                (individual.v_norm_sensor_activation).reshape((1, 7)))[0]

        scaled_output = np.array(
            [scale(xi, -2.0, 2.0) for xi in net_output])

        # Collect behavioral feature data
        wheel_speeds.append(net_output)
        sensor_activations.append(
            list(map(lambda x: 1 if x > 0.0 else 0, individual.v_norm_sensor_activation)))

        # set motor wheel speeds
        individual.v_set_motors(*list(scaled_output))

        # sample every 2 second robot position
        if round(runtime, 2) % 2.0 == 0.0:
            # get vrep robot current position
            # print(individual.v_get_position())
            position.append(individual.v_get_position())

        # After this call, the first simulation step is finished
        # Now we can safely read all  values
        vrep.simxGetPingTime(individual.client_id)

        runtime += dt
        steps += 1

        # every 10 seconds the robot is in the same position given a threshold stop the simulation
        if round(runtime, 2) % 5.0 == 0.0:
            # print(euclidean_distance(position[0], position[-1]))
            if (euclidean_distance(position[0], position[-1])) < .05:
                collision = True

        #  fitness_t at time stamp
        (
            fitness_t,
            wheel_center,
            straight_movements,
            obstacles_distance
        ) = f_t_obstacle_avoidance(
            scaled_output, individual.v_norm_sensor_activation, 'vrep')

        fitness_agg = np.append(fitness_agg, fitness_t)

        # dump individuals data
        if settings.debug:
            save_debug_data(
                settings.path,
                individual.id,
                individual.v_sensor_activation,
                individual.v_norm_sensor_activation,
                net_output,
                scaled_output,
                wheel_center,
                straight_movements,
                obstacles_distance,
                fitness_t,
                'VREP',
                position[-1]
            )

    # calculate the fitnesss
    fitness = np.sum(fitness_agg)/settings.run_time

    # Now send some data to V-REP in a non-blocking fashion:
    vrep.simxAddStatusbarMessage(
        individual.client_id, 'genome_id: {} fitness: {:.4f} runtime: {:.2f} s'.format(
            individual.id, fitness, runtime), vrep.simx_opmode_oneshot)

    # Before closing the connection to V-REP, make sure that the last command sent out had time to arrive. You can guarantee this with (for example):
    vrep.simxGetPingTime(individual.client_id)

    if (vrep.simxStopSimulation(individual.client_id, settings.op_mode) == -1):
        return

    if type(genome).__name__ == 'Individual':
        generation = genome.gen

    behavioral_features = calc_behavioral_features(
        areas_counter,
        wheel_speeds,
        sensor_activations,
        settings.path,
        genome.key,
        generation,
        'VREP'
    )

    print('genome_id: {} fitness: {:.4f} runtime: {:.2f} s steps: {}'.format(
        genome.key, fitness, runtime, steps))

    if settings.debug:
        print('behavioral_features: {0}\n pos_sample: {1}\n n_samples: {2}\n'.format(
            behavioral_features, position, len(position)))

    if type(genome).__name__ == 'Individual':
        genome.features = np.array(behavioral_features)
        genome.task_fitness = fitness
        genome.position = np.array(position)
        genome.evaluation = 'VREP'
        genome.weights = network.get_weights()
        genome.sim_time = runtime
        # Save the neural network model
        model.save(settings.path + 'keras_models/' +
                   str(individual.id) + '_model.h5')

    if type(genome).__name__ == 'DefaultGenome':
        genome.features = behavioral_features
        genome.fitness = fitness
        genome.position = position

    time.sleep(1)

    return fitness

#####################################################
#               AUXILIARY METHODS                   #
#               AUXILIARY METHODS                   #
#               AUXILIARY METHODS                   #
#               AUXILIARY METHODS                   #
#####################################################

def thymio_get_position_every_2s(position):
    robot_m = get_marker_object(7)
    if robot_m.realxy() is not None:
        position.append(robot_m.realxy()[:2])

def init_network(genome, config=None, model=None):

    network = None

    if type(genome).__name__ == 'Individual':
        if True:
            # TODO the sizeing is hardcoded.
            # 7 NEURONS
            # weights = [
            #     np.array(genome[:35]).reshape(genome.weights_shape[0]),
            #     np.array(genome[35:40]).reshape(genome.weights_shape[1]),
            #     np.array(genome[40:50]).reshape(genome.weights_shape[2]),
            #     np.array(genome[-2:]).reshape(genome.weights_shape[3]),
            # ]
            # 14 NEURONS
            weights = [
                np.array(genome[:98]).reshape(genome.weights_shape[0]),
                np.array(genome[98:112]).reshape(genome.weights_shape[1]),
                np.array(genome[112:140]).reshape(genome.weights_shape[2]),
                np.array(genome[-2:]).reshape(genome.weights_shape[3]),
            ]
            # 2 HIDDEN LAYERS AND 14 NEURONS
            # weights = [
            #     np.array(genome[:98]).reshape(genome.weights_shape[0]),
            #     np.array(genome[98:112]).reshape(genome.weights_shape[1]),
            #     np.array(genome[112:308]).reshape(genome.weights_shape[2]),
            #     np.array(genome[308:322]).reshape(genome.weights_shape[3]),
            #     np.array(genome[322:350]).reshape(genome.weights_shape[4]),
            #     np.array(genome[-2:]).reshape(genome.weights_shape[5]),
            # ]
            # 10 NEURONS
            # weights = [
            #     np.array(genome[:70]).reshape(genome.weights_shape[0]),
            #     np.array(genome[70:80]).reshape(genome.weights_shape[1]),
            #     np.array(genome[80:100]).reshape(genome.weights_shape[2]),
            #     np.array(genome[-2:]).reshape(genome.weights_shape[3]),
            # ]

            model.set_weights(weights)
        else:
            model.set_weights(genome)

        network = model

    if type(genome).__name__ == 'DefaultGenome':
        network = neat.nn.FeedForwardNetwork.create(genome, config)

    return network

def build_scene(config_scene, settings, client_id):
    if config_scene:
            # Get the position of all the obstacles in reality
        obstacles_pos = [get_marker_object(obstacle).realxy()
                         for obstacle in (9, 10, 11)]
        # Get all obstacle handlers from VREP
        obstacle_handlers = [get_object_handle(client_id, obstacle) for obstacle in (
            'obstacle', 'obstacle1', 'obstacle0')]
        # Set the position of obstacles in vrep according the obstacles from reality
        for obs, handler in zip(obstacles_pos, obstacle_handlers):
            set_pose(client_id, handler, [obs[0], obs[1], 0.099999])

        # add markers position to obstacle_markers
        for position, marker in zip(obstacles_pos, settings.obstacle_markers):
            for _, value in marker.items():
                value.update(center=(position[:2]*1000).astype(int))

        obstacle_grid = create_grid(settings.obstacle_markers)
    else:
        # add markers position to obstacle_markers
        obstacles_pos = [[620, 590, 0], [880, 100, 0], [150, 430, 0]]
        for position, marker in zip(obstacles_pos, settings.obstacle_markers):
            for _, value in marker.items():
                value.update(center=position[:2])
        obstacle_grid = create_grid(settings.obstacle_markers)

    return obstacle_grid

def init_collision(client_id):
    # collistion detection initialization
    _, collision_handle = vrep.simxGetCollisionHandle(
        client_id, 'wall_collision', vrep.simx_opmode_blocking)
    _, collision = vrep.simxReadCollision(
        client_id, collision_handle, vrep.simx_opmode_streaming)

    return collision_handle, collision


import pickle
import os
import time
import numpy as np
import sys
import shutil
from openmmlib import openmmlib
from openmmlib import polymerutils
from openmmlib.polymerutils import scanBlocks
from openmmlib.openmmlib import Simulation
from openmmlib.polymerutils import grow_rw
from looplib import looptools
import pyximport; pyximport.install()
from smcTranslocator_diffusive_mixed.pyx import smcTranslocatorDirectional

'''
Make sure that on average we have 1 smcstep/200 3D simulations
eg by setting
smcStepsPerBLock=1
steps=200

if smcStepsPerBlock<1, reduce the number of 3D simulations per simulation block accordingly
The speed of the smc is 1-PAUSEP.

One consideration is that the number of smc steps between 3D simulation blocks should never exceed 1 too much, as it will result in jerky motion of the polymer.
'''

# -------defining parameters----------
# -- basic loop extrusion parameters--
GPU = sys.argv[1]
LIFETIME = int(sys.argv[2])   # Processivity
SEPARATION = int(sys.argv[3]) # Separation LEFs in number of monomers

PAUSEP=float(sys.argv[4]) # pause prob active arm, set to 0 if not using diffusive simulations
SLIDE_PAUSEP=float(sys.argv[5]) # pause prob passive arm

TADSizes =  [400,100,200,400,100,200,400,100,200,400,100,200,400,100,200,400,100,200,400,100,200,400,100,200]
N = sum(TADSizes) # number of monomers
smcStepsPerBlock = 1# I take something like 1/steppingprobability, because stepping is not determistic. I usually choose the probability of stepping to be max 0.1.
stiff = 0                 # Polymer siffness in unit of bead size
dens = 0.2 # density in beads / volume. The density can roughly be estimated by looking at the amount of DNA in a nuclear volume.
box = (N / dens) ** 0.33  # Define size of the bounding box for Periodic Boundary Conditions
data = polymerutils.grow_rw(N, int(box) - 2)  # creates a compact conformation
block = 0  # starting block

dt = 1*(1-PAUSEP)# Time step of each smcTranslocatorPol iteration, the time step is chosen such that MaxRate*dt=0.1,
#as this should give proper time step distributions.
steps =  int(200*dt) # nr of 3D simulation blocks btw advancing LEFs. For deterministic stepping choose 200-250 steps per block, otherwise, rescale with stepping probability.
stg = 0.8 #Probability of stalling at TAD boundary

BELT_ON=0
BELT_OFF=1
SWITCH=0
PUSH=0
PAIRED=0
SLIDE=1# diffusive motion
loop_prefactor=1.5  # Should be the same as in smcTranslocator
FULL_LOOP_ENTROPY=1 # Should be the same as in smcTranslocator
FRACTION_ONESIDED=float(sys.argv[6])

if dt==0:
    syst.exit('WARNING: dt=0')

'!!!!!!!Uncomment if entropy is considered'
if (SLIDE_PAUSEP < 1-np.exp(0.5*loop_prefactor)) or (SLIDE_PAUSEP < 1-0.5/np.cosh(0.5*loop_prefactor)):
    sys.exit('WARNING: slidepause prob is too small! simulations will bias toward forward sliding!')

#if ((PAUSEP>0.9) and (smcStepsPerBlock<10)):
#    sys.exit('WARNING: Make sure that smc steps per block is effectively 1')

saveEveryBlocks =  int(100/(smcStepsPerBlock*dt))#int(400/(smcStepsPerBlock*dt))   #  10 save every 10 blocks (saving every block is now too much almost)
skipSavedBlocksBeginning = int(20/(smcStepsPerBlock*dt))  # how many blocks (saved) to skip after you restart LEF positions
totalSavedBlocks = 4000#5000  # how many blocks to save (number of blocks done is totalSavedBlocks * saveEveryBlocks)
restartMilkerEveryBlocks = int(100/(smcStepsPerBlock*dt))   #int(100/(smcStepsPerBlock*dt))
#Only one Hamiltonian can be loaded at a time to the simkt, but we want to update the bonds every time a LEF steps. Loading a new Hamiltonian costs a lot of time. Instead we precalculate bonds and load all positions at once as one big Hamiltonian and just change the prefactors.

# parameters for smc bonds
smcBondWiggleDist = 0.2
smcBondDist = 0.5

if len(sys.argv)!=7:
    sys.exit('Number of input arguments is not correct')

#folder and number of blocks
#folder = "trajectory"
folder = "PolymerSimulations"

FullFileName=os.path.join(folder,"loopextrusion_Lifetime={0}_separation={1}_density={2}_N={3}_saveEveryBlocks={4}_totalSavedBlocks={5}_SLIDE_PAUSEPROB={6}_PAUSEPROB={7}_dt={8}_FRACTION_ONESIDED={9}_entropy".format(LIFETIME, SEPARATION, dens,N,saveEveryBlocks,totalSavedBlocks,SLIDE_PAUSEP,PAUSEP,dt,FRACTION_ONESIDED))

#FullFileName=os.path.join(folder,"NestedLoopTest2")

if os.path.isdir(FullFileName):
    var=input('Folder already exists, press 9 to stop or any other key to empty and recreate the folder and continue')
    if eval(var)==9:
        sys.exit('Stopped simulation')
    else:
        shutil.rmtree(FullFileName) # Remove folder and content

# assertions for easy managing code below

assert restartMilkerEveryBlocks % saveEveryBlocks == 0
assert (skipSavedBlocksBeginning * saveEveryBlocks) % restartMilkerEveryBlocks == 0
assert (totalSavedBlocks * saveEveryBlocks) % restartMilkerEveryBlocks == 0
assert smcStepsPerBlock*(1-PAUSEP)<3 # max number of steps per smc block should not be too large to prevent 'jerky' polymer motion

savesPerMilker = restartMilkerEveryBlocks // saveEveryBlocks
milkerInitsSkip = saveEveryBlocks * skipSavedBlocksBeginning  // restartMilkerEveryBlocks
milkerInitsTotal  = (totalSavedBlocks + skipSavedBlocksBeginning) * saveEveryBlocks // restartMilkerEveryBlocks
print("Time step = {0}, Milker will be initialized {1} times, first {2} will be skipped".format(dt,milkerInitsTotal, milkerInitsSkip))

# create filenames for Ekin, Epot and time

Ekin_fname = os.path.join(FullFileName,'Ekin.txt')
Epot_fname = os.path.join(FullFileName,'Epot.txt')
time_fname = os.path.join(FullFileName,'time.txt')
Rg_fname   = os.path.join(FullFileName,'Rg.txt')
Par_fname  = os.path.join(FullFileName,'Pars.txt')

def save_Es_ts_Rg():
    with open(time_fname, "a+") as time_file:
        time_file.write('%f\n'%(a.state.getTime()/openmmlib.ps))
    with open(Ekin_fname, "a+") as Ekin_file:
        Ekin_file.write('%f\n'%((a.state.getKineticEnergy())/a.N/a.kT))
    with open(Epot_fname, "a+") as Epot_file:
        Epot_file.write('%f\n'%((a.state.getPotentialEnergy()) /a.N /a.kT))
    with open(Rg_fname, "a+") as Rg_file:
        Rg_file.write('%f\n'%(analysis_plot_lib.rg(a.getData())))

class smcTranslocatorMilker(object):

    def __init__(self, smcTransObject):
        """
        :param smcTransObject: smc translocator object to work with
        """
        self.smcObject = smcTransObject
        self.allBonds = []

    def setParams(self, activeParamDict, inactiveParamDict):
        """
        A method to set parameters for bonds.
        It is a separate method because you may want to have a Simulation object already existing

        :param activeParamDict: a dict (argument:value) of addBond arguments for active bonds
        :param inactiveParamDict:  a dict (argument:value) of addBond arguments for inactive bonds

        """
        self.activeParamDict = activeParamDict
        self.inactiveParamDict = inactiveParamDict

    def setup(self, bondForce,  blocks = 100, smcStepsPerBlock = 1):
        """
        A method that milks smcTranslocator object
        and creates a set of unique bonds, etc.

        :param bondForce: a bondforce object (new after simulation restart!)
        :param blocks: number of blocks to precalculate
        :param smcStepsPerBlock: number of smcTranslocator steps per block
        :return:
        """

        if len(self.allBonds) != 0:
            raise ValueError("Not all bonds were used; {0} sets left".format(len(self.allBonds)))

        self.bondForce = bondForce

        #precalculating all bonds
        allBonds = []
        for dummy in range(blocks):
            self.smcObject.steps(smcStepsPerBlock)
            left, right = self.smcObject.getSMCs()
            bonds = [(int(i), int(j)) for i,j in zip(left, right)]
            allBonds.append(bonds)

        self.allBonds = allBonds
        self.uniqueBonds = list(set(sum(allBonds, []))) # 'sum' preserves order and makes one long list with bonds, 'set' creates a set with left bonds from different time points ordered from small to large and eliminates two equal bonds (also if they were created by different LEFs at different times). List turns set into a list with unique bonds at different time points.

        # adding forces and getting bond indices
        self.bondInds = []
        self.curBonds = allBonds.pop(0) # pop(0) removes and returns first list of bonds

        for bond in self.uniqueBonds:
            paramset = self.activeParamDict if (bond in self.curBonds) else self.inactiveParamDict
            ind = bondForce.addBond(bond[0], bond[1], **paramset)
            self.bondInds.append(ind)
        self.bondToInd = {i:j for i,j in zip(self.uniqueBonds, self.bondInds)}
        return self.curBonds,[]

    def step(self, context, verbose=False):
        """
        Update the bonds to the next step.
        It sets bonds for you automatically!
        :param context:  context
        :return: (current bonds, previous step bonds); just for reference
        """
        if len(self.allBonds) == 0:
            raise ValueError("No bonds left to run; you should restart simulation and run setup  again")

        pastBonds = self.curBonds
        self.curBonds = self.allBonds.pop(0)  # getting current bonds
        bondsRemove = [i for i in pastBonds if i not in self.curBonds]
        bondsAdd = [i for i in self.curBonds if i not in pastBonds]
        bondsStay = [i for i in pastBonds if i in self.curBonds]
        if verbose:
            print("{0} bonds stay, {1} new bonds, {2} bonds removed".format(len(bondsStay),
                                                                            len(bondsAdd), len(bondsRemove)))
        bondsToChange = bondsAdd + bondsRemove
        bondsIsAdd = [True] * len(bondsAdd) + [False] * len(bondsRemove)
        for bond, isAdd in zip(bondsToChange, bondsIsAdd):
            ind = self.bondToInd[bond]
            paramset = self.activeParamDict if isAdd else self.inactiveParamDict
            self.bondForce.setBondParameters(ind, bond[0], bond[1], **paramset)  # actually updating bonds
        self.bondForce.updateParametersInContext(context)  # now run this to update things in the context
        return self.curBonds, pastBonds

def initModel():

    birthArray = np.ones(N, dtype=np.double)*0.1
    deathArray = np.zeros(N, dtype=np.double) + 1. / (LIFETIME/dt)
    stallLeftArray = np.zeros(N, dtype=np.double)#Probability of stalling. Choose np.zeros(N, dtype=np.double) if no stalling
    stallRightArray = np.zeros(N, dtype=np.double)

    pauseArray = np.zeros(N, dtype=np.double) + PAUSEP
    slidePauseArray = np.zeros(N, dtype=np.double) + SLIDE_PAUSEP

    stallDeathArray = np.zeros(N, dtype=np.double) + 1 / (LIFETIME/dt)
    smcNum = N // SEPARATION
    curPos = 0
    for size in TADSizes: # setting positions & strength of boundary elements, looping over each TAD
        curPos += size
        if curPos < len(stallLeftArray):
            stallLeftArray[curPos] = stg
            stallRightArray[curPos] = stg

    oneSidedArray = np.ones(smcNum, dtype=np.int64)
    for i in range(int((1.-FRACTION_ONESIDED)*smcNum)):
        oneSidedArray[i] = 0

    belt_on_array = np.zeros(smcNum, dtype=np.double) + BELT_ON
    belt_off_array = np.zeros(smcNum, dtype=np.double) + BELT_OFF

    spf=slidePauseArray*(1.-(1.-SLIDE_PAUSEP)*np.exp(-1.*loop_prefactor))
    spb=slidePauseArray*(1.-(1.-SLIDE_PAUSEP)*np.exp(loop_prefactor))
    #sps=spf+spb
    SMCTran = smcTranslocatorDirectional(birthArray, deathArray, stallLeftArray, stallRightArray, pauseArray,
                                         stallDeathArray, smcNum, oneSidedArray, paired=PAIRED, slide=SLIDE,
                                         slidepauseForward=spf, slidepauseBackward=spb,
                                         #slidepauseSum=sps,
                                         switch=SWITCH, pushing=PUSH, belt_on=belt_on_array, belt_off=belt_off_array,SLIDE_PAUSEPROB=SLIDE_PAUSEP)

    return SMCTran

SMCTran = initModel()  # defining actual smc translocator object

# now polymer simulation code starts

# ------------feed smcTran to the milker---
SMCTran.steps(1000000)  # first steps to "equilibrate" SMC dynamics. If desired of course.
milker = smcTranslocatorMilker(SMCTran)   # now feed this thing to milker (do it once!)
#--------- end new code ------------

for milkerCount in range(milkerInitsTotal):
    doSave = milkerCount >= milkerInitsSkip

    # simulation parameters are defined below
    a = Simulation(timestep=80, thermostat=0.01)#Collision rate in inverse picoseconds, low collistion rate means ballistic like motion, default in openmmpolymer is 0.001. Motion polymer is not diffusive, this is ok for statistical average,
    #but not for dynamics of the polymer
    a.setup(platform="CUDA", PBC=True, PBCbox=[box, box, box], GPU=GPU, precision="mixed")  # set up GPU here, PBC=Periodic Boundary Conditions. Default integrator is langevin with 300 K, friction coefficient of 1/ps, step size 0.002ps
    a.saveFolder(FullFileName)
    a.load(data)
    a.addHarmonicPolymerBonds(wiggleDist=0.1) # WiggleDist controls distance at which energy of bond equals kT
    if stiff > 0:
        a.addGrosbergStiffness(stiff) # Chain stiffness is introduced by an angular potential U(theta)=stiff(1-cos(theta-Pi))
    a.addPolynomialRepulsiveForce(trunc=1.5, radiusMult=1.05) #Polynomial repulsive potential between particles. Has value trunc=3.0 at zero, stays flat until 0.6-0.7 and then drops to zero. For attraction between a selective set of particles, use LeonardJones or addSelectiveSSWForce (see blocks.py or ask Johannes)
    a.step = block

    # ------------ initializing milker; adding bonds ---------
    # copied from addBond
    kbond = a.kbondScalingFactor / (smcBondWiggleDist ** 2)
    bondDist = smcBondDist * a.length_scale

    activeParams = {"length":bondDist,"k":kbond}
    inactiveParams = {"length":bondDist, "k":0}
    milker.setParams(activeParams, inactiveParams)

    # this step actually puts all bonds in and sets first bonds to be what they should be
    milker.setup(bondForce=a.forceDict["HarmonicBondForce"],
                blocks=restartMilkerEveryBlocks,   # default value; milk for 100 blocks
                 smcStepsPerBlock=smcStepsPerBlock)  #
    print("Restarting milker")

    a.doBlock(steps=steps, increment=False)  # do block for the first time with first set of bonds in
    for i in range(restartMilkerEveryBlocks - 1):
        curBonds, pastBonds = milker.step(a.context)  # this updates bonds. You can do something with bonds here
        if i % saveEveryBlocks == (saveEveryBlocks - 2):
            a.doBlock(steps=steps, increment = doSave)
            if doSave:
                a.save()
                pickle.dump(curBonds, open(os.path.join(a.folder, "SMC{0}.dat".format(a.step)),'wb'))
                save_Es_ts_Rg() # save energies and time
        else:
            a.integrator.step(steps)  # do steps without getting the positions from the GPU (faster)

    data = a.getData()  # save data and step, and delete the simulation
    block = a.step
    del a

    time.sleep(0.2)  # wait 200ms for sanity (to let garbage collector do its magic)

with open(Par_fname,"a+") as Parfile:
    Parfile.write(" tau="+str(LIFETIME)+"\n Separation="+str(SEPARATION)+"\n N="+str(N)+"\n smcStepsPerBlock="+str(smcStepsPerBlock)+"\n stiff="+str(stiff)+"\n dens="+str(dens)+"\n block="+str(block) + "\n SaveEveryBlocks="+str(saveEveryBlocks)+"\n skipSavedBlocksBeginning="+str(skipSavedBlocksBeginning)+"\n totalSavedBlocks="+str(totalSavedBlocks)+"\n restartMilkerEveryBlocks="+str(restartMilkerEveryBlocks)+"\n smcBondWiggleDist="+str(smcBondWiggleDist)+"\n smcBondDist="+str(smcBondDist)+"\n SmcTimestep=" + str(dt)+"\n Mean TADsize"+str(np.mean(TADSizes)))


import numpy as np
import cv2
import dlib
from PIL import Image
import grpc

# type of modes to apply face swapping
ALL_FACE_MODE = 'apply_on_all'
LARGEST_FACE_MODE = 'choose_largest_face'

# type of image arguments for referring image in error messages
ARG_IMG_SRC = 'face'
ARG_IMG_DST = 'meme'

# model file name for face landmark detection in dlib library
FACE_LANDMARK_SHAPE_DETECTOR_FILENAME = 'shape_predictor_68_face_landmarks.dat'

NUMBER_OF_FACE_LANDMARKS = 68   # number of facelandmark points

class GRPCException(Exception):
    def __init__(self, code, message):
        self.code = code
        self.message = message

    def __str__(self):
        return "GRPC Exception {}: {}".format(self.code, self.message)

def shape_to_np_array(shape, points_num=NUMBER_OF_FACE_LANDMARKS, dtype="int32"):
    """
    Change shape data structure to np array

    Args:
        shape: dlib shape data type
        dtype: numpy data type

    Returns:
        np_array from dlib shape data type
    """
    # initialize our nparray
    np_array = np.zeros((points_num, 2), dtype=dtype)

    # loop over all points and convert them to (x, y) tuples
    for i in range(points_num):
        np_array[i] = (shape.part(i).x, shape.part(i).y)

    return np_array

def is_black_and_white(img):
    img_arr = np.array(img)
    channel_colors_match = (img_arr[:, :, 0] == img_arr[:, :, 1]) == (
        img_arr[:, :, 1] == img_arr[:, :, 2])
    if channel_colors_match.all() == True:
        return True
    return False

def match_image_color(src_img, dst_img):
    if is_black_and_white(src_img) is True and is_black_and_white(dst_img) is False:
        dst_img_gray = cv2.cvtColor(dst_img, cv2.COLOR_BGR2GRAY)
        dst_img_bw = cv2.cvtColor(dst_img_gray, cv2.COLOR_GRAY2BGR)
        dst_img = dst_img_bw
    elif is_black_and_white(src_img) is False and is_black_and_white(dst_img) is True:
        src_img_gray = cv2.cvtColor(src_img, cv2.COLOR_BGR2GRAY)
        src_img_bw = cv2.cvtColor(src_img_gray, cv2.COLOR_GRAY2BGR)
        src_img = src_img_bw
    else:
        pass
    return src_img, dst_img

def get_face_landmark_points(img, arg_img_type):
    """
    Generate facial landmark points of a given image

    Args:
        img: nparray
            Readed value of an image
        arg_image_type: str from { 'face', 'meme' }
            wherther img is input face image or meme image for error message purposes

    Returns:
        list of nparrays:
            list of landmark points for every face in a given image
    """

    # convert the image to greyscaleprint("land here")
    try:
        img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    except Exception:
        raise GRPCException(
            "{} image could not be converted to grayscale. "
            "Please, check if appropriate image is selected".format(arg_img_type))

    # detect the face then find the landmarks points
    detector = dlib.get_frontal_face_detector()
    try:
        predictor = dlib.shape_predictor(
            FACE_LANDMARK_SHAPE_DETECTOR_FILENAME)
    except Exception:
        raise GRPCException(grpc.StatusCode.FAILED_PRECONDITION,
                            "Facial landmark points file cann't be found. "
                            "Download 'shape_predictor_68_face_landmarks.dat'")

    faces = detector(img_gray)
    if len(faces) == 0:
        raise GRPCException(grpc.StatusCode.FAILED_PRECONDITION,
                            'No face could be detected from {} image input.'.format(arg_img_type))

    faces_landmark_points = []
    for face in faces:
        landmarks = predictor(img_gray, face)
        landmark_points = shape_to_np_array(landmarks)
        faces_landmark_points.append(landmark_points)

    return faces_landmark_points

def rectContains(rect, point):
    """Check if a point is inside a rectangle.

    Parameters
    ----------
    rect : tuple
        Points of the rectangle edges.
    point : tuple
        List of points.

    Returns
    -------
    bool
        Return true if the points are inside rectangle else return false.

    """
    if point[0] < rect[0]:
        return False
    elif point[1] < rect[1]:
        return False
    elif point[0] > rect[0] + rect[2]:
        return False
    elif point[1] > rect[1] + rect[3]:
        return False
    return True

def is_overlap(Rect1, Rect2):
    """Check if two rectangle are overlapped.

    Parameters
    ----------
    Rect1 : tuple
        Points of the rectangle1 edges.
    Rect2 : tuple
        Points of the rectangle2 edges.

    Returns
    -------
    bool
        Return true is overlap else return false.

    """
    if rectContains(Rect1, Rect2) or rectContains(Rect2, Rect1):
        return True
    else:
        return False

def fix_landmark_overlap(faces_landmark_points):
    """Disregared the overlapped faces(one of them).

    Parameters
    ----------
    faces_landmark_points : list
        Landmark points of each faces.

    Returns
    -------
    bool
        return the non overlapped faces landmark points.

    """
    none_overlap_faces_landmark_points = []
    faces = len(faces_landmark_points)
    overlapped_faces = []  # store the overlapped faces
    for i in range(faces):
        for j in range(faces):
            # check if face[i] and face[j] are the same face and are not overlapped with any other faces
            if (i == j) and (j not in overlapped_faces):
                none_overlap_faces_landmark_points.append(
                    faces_landmark_points[i])
            # check if the bounding box are overlapped
            elif is_overlap(cv2.boundingRect(faces_landmark_points[i]), cv2.boundingRect(faces_landmark_points[j])):
                overlapped_faces.append(j)
                continue
    return none_overlap_faces_landmark_points

def choose_largest_face(faces_landmark_points):
    """
    Choose largest face from all the faces in a given image

    Args:
        faces_landmark_points : list of nparrays
            landmark points of all the faces

    Returns:
        nparray:
            landmark points of the largest face
    """
    faces_num = len(faces_landmark_points)

    # if only one face return it
    if faces_num == 1:
        return faces_landmark_points[0]

    # search for the largest face otherwise
    largest_size = 0
    for face_i in range(faces_num):
        x, y, w, h = cv2.boundingRect(faces_landmark_points[face_i])

        face_size = w * h   # ? getting face size
        if face_size > largest_size:
            largest_size = face_size
            largest_face_index = face_i

    largest_face_landmark_points = faces_landmark_points[largest_face_index]

    return largest_face_landmark_points

def find_face_direction(landmark_points):
    """
    Find left or right direction of a face.

    Args:
        landmark_points : nparray
            Facial lanmark points of a given image

    Returns:
        direction: str
            Direction of the face
    """

    # ? are these points enough
    pt1 = landmark_points[3]
    pt2 = landmark_points[34]
    pt3 = landmark_points[15]

    # ? are these criteria enough
    face_width = np.linalg.norm(np.subtract(pt3, pt1))
    left_dist = np.linalg.norm(np.subtract(pt2, pt1)) / face_width
    right_dist = np.linalg.norm(np.subtract(pt3, pt2)) / face_width

    # ? emprical constant 0.2
    if left_dist > right_dist + 0.2:
        direction = "right"
    elif right_dist > left_dist + 0.2:
        direction = "left"
    else:
        direction = "front"  # ? what happens when front

    return direction

def align_face_direction(src_img, src_landmark_points, dst_landmark_points):
    """
    Flip src_img based on the alignment of two faces

    Args:
        src_img : nparray
            Numpy array of src_img.
        src_landmark_points : nparray
            Landmark points of a face from src_img
        dst_landmark_points : nparray
            Landmark points of a face from dst_img

    Returns:
        nparray:
            The flipped or the original image
    """

    src_direction = find_face_direction(src_landmark_points)
    dst_direction = find_face_direction(dst_landmark_points)

    # ? how to align if front face detected
    if (src_direction == "left" and dst_direction == "right") or (src_direction == "right" and dst_direction == "left"):
        flipped_src_img = cv2.flip(src_img, flipCode=1)
        return flipped_src_img

    return src_img

def get_convex_hulls(src_face_points, dst_face_points):
    """
    Find the convex hulls for both landmark points

    Args:
        src_face_points : nparray
            face landmark points from src_image
        dst_face_points : nparray
            face landmark points from dst_image

    Returns:
        list:
            convex hull points bounding src face points
        list:
            convex hull points bounding dst face points
    """

    # Find convex hull of the two landmark points
    src_hull_points = []
    dst_hull_points = []

    src_hull_index = cv2.convexHull(src_face_points, returnPoints=False)
    dst_hull_index = cv2.convexHull(dst_face_points, returnPoints=False)

    # use the max number of convex hull points.
    # ? should we use larger or smaller(for the sake of only avoid index out of bound)
    if len(src_hull_index) > len(dst_hull_index):
        max_hull_index = src_hull_index
    else:
        max_hull_index = dst_hull_index

    for i in range(len(max_hull_index)):
        src_hull_points.append(src_face_points[int(max_hull_index[i])])
        dst_hull_points.append(dst_face_points[int(max_hull_index[i])])

    return src_hull_points, dst_hull_points

def get_delaunay_triangles(img, points):
    """
    Calculate delauney triangles of a given points

    Args:
        img : nparray
            image
        points : list
            points to form delaunay triangles

    Returns:
        list:
            triple indexes of points that form delanauy triangles
    """
    # ? should we calculate all delanuay triangles
    (h, w, _) = img.shape
    rect = (0, 0, w, h)

    subdiv = cv2.Subdiv2D(rect)
    subdiv.insert(points)
    triangleList = subdiv.getTriangleList()

    delaunay_triangle_indexes = []

    for t in triangleList:
        vtx1 = (t[0], t[1])
        vtx2 = (t[2], t[3])
        vtx3 = (t[4], t[5])
        vertexes = [vtx1, vtx2, vtx3]

        point_indexes = []
        # Get face-points (from 68 face detector) by coordinates
        for i in range(3):
            for j in range(0, len(points)):
                if(vertexes[i][0] == points[j][0] and vertexes[i][1] == points[j][1]):
                    point_indexes.append(j)
                    break
        # Three points form a triangle
        assert(len(point_indexes) == 3)
        delaunay_triangle_indexes.append(
            (point_indexes[0], point_indexes[1], point_indexes[2]))

    return delaunay_triangle_indexes

def replace_triangle(src_img, dst_img, src_tri_points, dst_tri_points, DEBUG=False):
    """
    Warp src_tri to dst_tri;then replace dst_tri portion of dst_img by src_tri portion of src_img
    Changes dst_img by replacing the warped triangle

    Args:
        src_img: numpy.ndarray
            src_img
        dst_img : numpy.ndarray
            dst_img
        src_tri_points : list of triple points
            a single triangle points from src_img
        dst_tri_points : list of triple points
            a single triangle points of dst_img
    """

    # Find bounding rectangle for each triangle
    src_tri_roi = cv2.boundingRect(np.array(src_tri_points))
    src_roi_x, src_roi_y, src_roi_w, src_roi_h = src_tri_roi

    dst_tri_roi = cv2.boundingRect(np.array(dst_tri_points))
    dst_roi_x, dst_roi_y, dst_roi_w, dst_roi_h = dst_tri_roi

    # Transform points by left top corner of the respective rectangles
    transformed_src_tri = []
    transformed_dst_tri = []

    if DEBUG:
        dst_tri_vtxs = np.array(dst_tri_points)
        dst_roi_vtxs = np.array([
            (dst_roi_x, dst_roi_y),
            (dst_roi_x+dst_roi_w, dst_roi_y),
            (dst_roi_x+dst_roi_w, dst_roi_y+dst_roi_h),
            (dst_roi_x, dst_roi_y+dst_roi_h)
        ])
        cv2.drawContours(dst_img, [dst_tri_vtxs], 0, (255, 0, 0), 1)
        # cv2.drawContours(dst_img, [rect_vtxs], 0, (0, 0, 255), 1)
    for vtx_i in range(3):
        if DEBUG:
            cv2.circle(dst_img, tuple(dst_tri_points[vtx_i]), 2, (255, 0, 0))

        transformed_src_tri.append(
            ((src_tri_points[vtx_i][0] - src_roi_x), (src_tri_points[vtx_i][1] - src_roi_y)))
        transformed_dst_tri.append(
            ((dst_tri_points[vtx_i][0] - dst_roi_x), (dst_tri_points[vtx_i][1] - dst_roi_y)))

    # Apply warpImage to small rectangular patches
    src_img_roi = src_img[src_roi_y:src_roi_y +
                          src_roi_h, src_roi_x:src_roi_x + src_roi_w]

    # find convertion matrix from src_tri to dst_tri
    warpMat = cv2.getAffineTransform(
        np.float32(transformed_src_tri), np.float32(transformed_dst_tri))

    # ? warping method
    dst_roi_dims = (dst_roi_w, dst_roi_h)
    warped_src_img_roi = cv2.warpAffine(src_img_roi, warpMat, dst_roi_dims, None,
                                        flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT)

    # Get mask by filling triangle
    warped_src_mask = np.zeros(
        warped_src_img_roi.shape, warped_src_img_roi.dtype)
    cv2.fillConvexPoly(warped_src_mask, np.int32(
        transformed_dst_tri), (1, 1, 1))
    warped_src_img_roi = warped_src_img_roi * warped_src_mask

    # Copy triangular region of the rectangular patch to the output image
    dst_img_roi = dst_img[dst_roi_y:dst_roi_y +
                          dst_roi_h, dst_roi_x:dst_roi_x + dst_roi_w]
    dst_roi_mask = (1, 1, 1) - warped_src_mask
    dst_img[dst_roi_y:dst_roi_y + dst_roi_h, dst_roi_x:dst_roi_x + dst_roi_w] \
        = dst_img_roi * dst_roi_mask + warped_src_img_roi

    if DEBUG:
        cv2.imshow('warped_src_img_roi', warped_src_img_roi)
        cv2.imshow('warped_src_mask', warped_src_mask)
        cv2.imshow('dst_img_roi', dst_img_roi)
        cv2.imshow('dst_img', dst_img)
        cv2.waitKey()

def replace_all_triangles(src_img, dst_img, dst_delaunay_triangle_indexes, src_points, dst_points):
    """
    Compute warp triangles for each triangles of image1 and image2.first find.
    corresponding landmark points of each triangles from the triangulations
    indexes. then warp each triangle of image1 to image2 by calling created
    warpTriangle function.
    Changes dst_img by replacing all delaunay triangles

    Args:
        src_img: numpy.ndarray
            src_img
        dst_img: numpy.ndarray
            dst_img
        dst_delaunay_triple_indexes : list
            triplewise indexes that form delaunay triangle of dst_points
        src_points : list
            src_img points corresponding to dst_img
        dst_points : list
            dst_img points corresponding to src_img
    """

    # iterate through each triangles
    for triangle_i in range(len(dst_delaunay_triangle_indexes)):
        src_tri_points = []
        dst_tri_points = []

        # iterate through all the three triangle indexes and find t1 and t2
        for vtx_j in range(3):
            src_tri_points.append(
                src_points[dst_delaunay_triangle_indexes[triangle_i][vtx_j]])
            dst_tri_points.append(
                dst_points[dst_delaunay_triangle_indexes[triangle_i][vtx_j]])

        replace_triangle(
            src_img, dst_img, src_tri_points,  dst_tri_points)

def apply_seamless_clone(src_img, dst_img, dst_points):
    """
    Seamlessly clone portion of src image and copy it to dst image

    Args:
        src : type
            Description of parameter `src`.
        dst : type
            Description of parameter `dst`.
        dst_points : list of ndarray
            points where src will be placed on dst

    Returns:
        numpy.nparray
            portion dst replaced by portion of src
    """

    # calculate mask
    mask = np.zeros(dst_img.shape, dtype=dst_img.dtype)
    cv2.fillConvexPoly(mask, np.int32(dst_points), (255, 255, 255))

    # calculate center dst image where center of src image put
    x, y, w, h = cv2.boundingRect(np.float32([dst_points]))
    dst_points_center = ((x + round(w / 2), y + round(h / 2)))

    # ? check MIXED_CLONE mode
    swappedImage = cv2.seamlessClone(
        src_img, dst_img, mask, dst_points_center, cv2.NORMAL_CLONE)

    return swappedImage

def swap_using_landmark_points(src_img, dynamic_dst_img, original_dst_img, src_face_points, dst_face_points):
    """
    This function contain code tha will be use for both mode. inorder to avoid repetition

    Args:
        src_img : nparray
            Readed value of source image
        dst_img : nparray
            Readed value of dst_img
        processed_dst_img : nparray
            Copy of dst image that may be changed by previous swaps
        src_face_points : nparray
            List of facial landmark poinst of single face from source image
        dst_face_points : nparray
            List of facial landmark poinst of single face from destination image

    Returns:
        numpy.ndarray
            The swapped image
    """
    # find the convex hull bounding the landmark points of the images
    src_hulls, dst_hulls = get_convex_hulls(src_face_points, dst_face_points)

    # calculate the delauney triangulations
    dst_delaunay_triangle_indexes = get_delaunay_triangles(
        dynamic_dst_img, dst_hulls)

    replace_all_triangles(
        src_img, dynamic_dst_img, dst_delaunay_triangle_indexes, src_hulls, dst_hulls)

    swapped_image = apply_seamless_clone(
        dynamic_dst_img, original_dst_img, dst_hulls)

    return swapped_image

def faceSwap(src_img, dst_img, mode=LARGEST_FACE_MODE, showImages=False):
    """
    Swap a face from src_img to dst_img

    Args:
        src_img: nparray
            Readed value of source image to be swapped on destination image
        dst_img: nparray
            Readed value of destination image
        mode:
            How the faces should be swapped
        showImages : bool
            An optional argument whether or not to show the original images

    Returns:
        nparray:
            The swapped image
    """
    # match the color of the two image
    src_img, dst_img = match_image_color(src_img, dst_img)

    # save the original dst img
    original_dst_img = np.copy(dst_img)

    # find landmark points of the images
    all_src_faces_points = get_face_landmark_points(src_img, ARG_IMG_SRC)
    src_face_points = choose_largest_face(all_src_faces_points)

    dst_faces_points = get_face_landmark_points(dst_img, ARG_IMG_DST)

    # if only one face is gonna be replaced it is ultimately largest_face_mode
    dst_faces_num = len(dst_faces_points)
    if dst_faces_num == 1:
        mode = LARGEST_FACE_MODE
    else:   # fix if different faces overlap
        dst_faces_points = fix_landmark_overlap(dst_faces_points)

    if mode == ALL_FACE_MODE:
        for face_i in range(dst_faces_num):
            dst_face_i_points = dst_faces_points[face_i]

            swappedImage = swap_using_landmark_points(
                src_img, dst_img, original_dst_img, src_face_points, dst_face_i_points)
            original_dst_img = swappedImage
    else:
        # find landmark points of the largest face in an image
        large_dst_face_points = choose_largest_face(dst_faces_points)

        # match the direction of the two images
        src_img = align_face_direction(
            src_img, src_face_points, large_dst_face_points)

        # recompute lanmark points on the flipped image or new one.
        # ? don't need to recompute, flip points too
        all_src_face_points = get_face_landmark_points(src_img, ARG_IMG_SRC)
        src_face_points = choose_largest_face(all_src_face_points)

        swappedImage = swap_using_landmark_points(
            src_img, dst_img, original_dst_img, src_face_points, large_dst_face_points)

    if showImages == True:
        show_images(src_img, original_dst_img, swappedImage,
                    showOriginalImages=True)

    return swappedImage

def show_images(img1, img2, swappedImage, showOriginalImages=False):
    """
    Display image1, image2 and warped image.

    Parameters
    ----------
    img1 : numpy.nparray
        output of image1 read with opencv.
    img2 : numpy:nparray
        output of image2 read with opencv.
    swappedImage : numpy.nparray
        the swapped image or new value of image2.

    Returns
    -------
    deosn't return any value. it just display the images.

    """
    if showOriginalImages:
        cv2.imshow("image1", img1)
        cv2.imshow("image2", img2)
    cv2.imshow("swappedImage", swappedImage)

    cv2.waitKey(0)
    cv2.destroyAllWindows()

if __name__ == '__main__':
    image1 = "faceswap/images/kalise.jpg"
    image2 = "faceswap/images/black_and_white.jpg"

    src_img = cv2.imread(image1)
    dst_img = cv2.imread(image2)

    swappedImage = faceSwap(src_img, dst_img, mode="apply_on_all",
                            showImages=True)


# Adapted from https://github.com/biubug6/Pytorch_Retinaface
# Original license: MIT
import torch
import numpy as np
from .. import torch_utils
import typing
from .models.retinaface import RetinaFace
from ..box_utils import batched_decode
from .utils import decode_landm
from .config import cfg_mnet, cfg_re50
from .prior_box import PriorBox
from torch.hub import load_state_dict_from_url
from torchvision.ops import nms
from ..base import Detector
from ..build import DETECTOR_REGISTRY

class RetinaNetDetector(Detector):

    def __init__(
            self,
            model: str,
            *args,
            **kwargs):
        super().__init__(*args, **kwargs)
        if model == "mobilenet":
            cfg = cfg_mnet
            state_dict = load_state_dict_from_url(
                "https://folk.ntnu.no/haakohu/RetinaFace_mobilenet025.pth",
                map_location=torch_utils.get_device()
            )
        else:
            assert model == "resnet50"
            cfg = cfg_re50
            state_dict = load_state_dict_from_url(
                "https://folk.ntnu.no/haakohu/RetinaFace_ResNet50.pth",
                map_location=torch_utils.get_device()
            )
            state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
        net = RetinaFace(cfg=cfg)
        net.eval()
        net.load_state_dict(state_dict)
        self.cfg = cfg
        self.net = net.to(self.device)
        self.mean = np.array([104, 117, 123], dtype=np.float32)

    def batched_detect_with_landmarks(
            self, image: np.ndarray) -> typing.Tuple[np.ndarray, np.ndarray]:
        """Takes N images and performs and returns a set of bounding boxes as
            detections
        Args:
            image (np.ndarray): shape [N, height, width, 3]
        Returns:
            np.ndarray: shape [N, 5] with (xmin, ymin, xmax, ymax, score)
            np.ndarray: shape [N, 5, 2] with 5 landmarks with (x, y)
        """
        image = image.astype(np.float32) - self.mean
        image = np.moveaxis(image, -1, 1)
        image = torch.from_numpy(image).to(self.device)
        height, width = image.shape[2:]
        boxes, landms = self._detect(image, return_landmarks=True)
        scores = boxes[:, :, -1]
        boxes = boxes[:, :, :-1]
        final_output_box = []
        final_output_landmarks = []
        for i in range(len(boxes)):
            boxes_ = boxes[i]
            landms_ = landms[i]
            scores_ = scores[i]
            # Confidence thresholding
            keep_idx = scores_ >= self.confidence_threshold
            boxes_ = boxes_[keep_idx]
            scores_ = scores_[keep_idx]
            landms_ = landms_[keep_idx]
            # Non maxima suppression
            keep_idx = nms(
                boxes_, scores_, self.nms_iou_threshold)
            boxes_ = boxes_[keep_idx]
            scores_ = scores_[keep_idx]
            landms_ = landms_[keep_idx]
            # Scale boxes
            boxes_[:, [0, 2]] *= width
            boxes_[:, [1, 3]] *= height
            # Scale landmarks
            landms_ = landms_.cpu().numpy().reshape(-1, 5, 2)
            landms_[:, :, 0] *= width
            landms_[:, :, 1] *= height
            dets = torch.cat(
                (boxes_, scores_.view(-1, 1)), dim=1).cpu().numpy()
            final_output_box.append(dets)
            final_output_landmarks.append(landms_)
        return final_output_box, final_output_landmarks

    @torch.no_grad()
    def _detect(
            self, image: np.ndarray,
            return_landmarks=False) -> np.ndarray:
        """Batched detect
        Args:
            image (np.ndarray): shape [N, H, W, 3]
        Returns:
            boxes: list of length N with shape [num_boxes, 5] per element
        """
        image = image[:, [2, 1, 0]]
        loc, conf, landms = self.net(image)  # forward pass
        scores = conf[:, :, 1:]
        height, width = image.shape[2:]
        priorbox = PriorBox(
            self.cfg, image_size=(height, width))
        priors = priorbox.forward()
        priors = torch_utils.to_cuda(priors, self.device)
        prior_data = priors.data
        boxes = batched_decode(loc, prior_data, self.cfg['variance'])
        boxes = torch.cat((boxes, scores), dim=-1)
        if return_landmarks:
            landms = decode_landm(landms, prior_data, self.cfg['variance'])
            return boxes, landms
        return boxes

@DETECTOR_REGISTRY.register_module
class RetinaNetResNet50(RetinaNetDetector):

    def __init__(self, *args, **kwargs):
        super().__init__("resnet50", *args, **kwargs)

@DETECTOR_REGISTRY.register_module
class RetinaNetMobileNetV1(RetinaNetDetector):

    def __init__(self, *args, **kwargs):
        super().__init__("mobilenet", *args, **kwargs)


import matplotlib.pyplot as plt
import numpy as np
import pyautogui
from PIL import Image
from keras.models import load_model
from multiprocess import Pool
from pynput.keyboard import KeyCode, Listener

from config import *
from time import time

model = load_model('recognize_new')  # 已训练的CNN

# 获取截图
def screen_shoot(rows=ROWS, cols=COLS, region=BG_REGION):
    resize_shape = (32, 32)
    board = []
    back_ground = np.array(pyautogui.screenshot(region=region))
    for j in range(rows):
        board_r = []
        for i in range(cols):
            l, u, r, d = [int(o) for o in
                          (i * cols_size, j * row_size, i * cols_size + cols_size, j * row_size + row_size)]
            crop = back_ground[u:d + 1, l:r + 1]
            crop = Image.fromarray(crop).resize(resize_shape)  # 将数组转成图像，resize之后再转成数组
            crop = np.array(crop)
            board_r.append(crop)
        board.append(board_r)
    return board

# 使用cnn识别图像
def get_cnn_board(model, rows=ROWS, cols=COLS, region=BG_REGION):
    # 将所有格子层叠排列，一起输入到cnn，比一个一个输入cnn计算速度快很多
    board = screen_shoot(rows, cols, region)
    sq_board = []
    for i in range(len(board)):
        sq_board += board[i]
    sq_board = np.array(sq_board) / 255.

    sq_pred = np.argmax(model.predict(sq_board), axis=1)  # 获取cnn的预测值
    sq_pred = np.array([net_decoder[i] for i in sq_pred])  # 解码cnn的预测值
    sq_pred = sq_pred.reshape((rows, cols))
    return sq_pred

# 显示截图结果
def show_screen_shoot(board=None):
    if board is None:
        board = screen_shoot()
    b = np.zeros((ROWS * 32, COLS * 32, 3))
    for i in range(ROWS):
        for j in range(COLS):
            b[i * 32:i * 32 + 32, j * 32:j * 32 + 32] = board[i][j]
    b = b / 255
    print(b)
    plt.imshow(b)
    plt.show()

# 显示识别结果
def show_recognized(board):
    plt.imshow(board)
    plt.show()

# 使用pyautogui识别图像
def get_board(rows=ROWS, cols=COLS, bg_region=BG_REGION, imags=None):
    # 旧版本使用pyautogui 识别图片，虽然调用的是opencv的接口但耗时太长
    board = -2 * np.ones((rows, cols), dtype='i2')
    back_ground = pyautogui.screenshot(region=bg_region)
    for num in imags:
        num_locations = pyautogui.locateAll(imags[num], back_ground, confidence=0.99)
        for loc in num_locations:
            center = pyautogui.center(loc)
            board[int(center[1] / row_size), int(center[0] // cols_size)] = num
    return board

# 将识别出来的图片保存到对应文件夹，可以增加训练样本
def record_board(model, path=PATH):
    resize_shape = (32, 32)
    sq_board = []
    back_ground = pyautogui.screenshot(region=BG_REGION)
    back_ground = np.array(back_ground)
    crops = []
    for j in range(ROWS):
        crops_r = []
        for i in range(COLS):
            l, u, r, d = [int(o) for o in
                          (i * cols_size, j * row_size, i * cols_size + cols_size, j * row_size + row_size)]
            crop = back_ground[u:d + 1, l:r + 1]
            crops_r.append(Image.fromarray(crop))

            crop = Image.fromarray(crop).resize(resize_shape)
            crop = np.array(crop)
            sq_board.append(crop)
        crops.append(crops_r)
    sq_board = np.array(sq_board)
    sq_pred = np.argmax(model.predict(sq_board), axis=1)
    sq_pred = np.array([net_decoder[i] for i in sq_pred])
    sq_pred = sq_pred.reshape((ROWS, COLS))
    for i in range(ROWS):
        for j in range(COLS):
            save_path = path + '\\' + str(sq_pred[i][j]) + '\\' + str(sq_pred[i][j]) + '_' + str(
                int(time() % 10000)) + '.png'
            crops[i][j].save(save_path)

# AI算法
def deal_board(board):
    to_flag = set()
    to_click = set()
    not_sure = dict()
    pad_board = np.pad(board, 1)
    row, col = board.shape
    while 1:
        action_nums = to_flag.__len__() + to_click.__len__()
        # 第一步，标记可以标记的旗子，把位置保存到to_flag
        for i in range(1, row + 1):
            for j in range(1, col + 1):
                if pad_board[i, j] > 0:
                    n9 = pad_board[i - 1:i + 2, j - 1:j + 2]
                    unknown_mask = (n9 == -2)
                    maybe_bomb = (n9 == -2) | (n9 == -10) | (n9 == -1)
                    if maybe_bomb.sum() == pad_board[i, j]:
                        n9[unknown_mask] = -1
                        for iii in range(3):
                            for jjj in range(3):
                                if unknown_mask[iii, jjj]:
                                    to_flag.add((iii + i - 1, jjj + j - 1))

        # 第二步消去旗子 用-10表示消去的旗子
        for i in range(1, row + 1):
            for j in range(1, col + 1):
                if pad_board[i, j] == -1:
                    pad_board[i, j] = -10
                    for ii in range(3):
                        for jj in range(3):
                            if pad_board[ii + i - 1, jj + j - 1] > 0:
                                pad_board[ii + i - 1, jj + j - 1] -= 1

        # 第三步，寻找可翻开的格子
        for i in range(1, row + 1):
            for j in range(1, col + 1):
                if pad_board[i, j] == 0:
                    for ii in range(3):
                        for jj in range(3):
                            if pad_board[ii + i - 1, jj + j - 1] == -2:
                                to_click.add((ii + i - 1, jj + j - 1))
                                pad_board[ii + i - 1, jj + j - 1] = -23

        # 第四步，正常办法解决不了，开启推理模式
        if action_nums == to_flag.__len__() + to_click.__len__():
            # 对每个没有翻开的格子a，如果周围有翻开的格子b1，那么将b1周围的所有未翻开的格子list1和b1的大小保存到一个列表，
            #                    如果周围有翻开的格子b2，那么将b2周围的所有未翻开的格子list2和b2的大小保存到一个列表，...
            # 最后得到一个关联关系  a的位置：与a有关的未翻开的格子的位置及隐藏地雷数量，保存到not_sure中
            for i in range(1, row + 1):
                for j in range(1, col + 1):
                    if pad_board[i, j] > 0:
                        not_sure_set = set()
                        for ii in range(3):
                            for jj in range(3):
                                if pad_board[ii + i - 1, jj + j - 1] == -2:
                                    not_sure_set.add((ii + i - 1, jj + j - 1))
                        for co_pos in not_sure_set:
                            not_sure[co_pos] = not_sure.get(co_pos, []) + ([[not_sure_set, pad_board[i, j]]])
            #        * * *             *   *  *   *   *
            #        * 1 ？            ？   4  ？  1   ？
            #        * 1 ？            ？   *  ？  *   ？
            #        * * ？            ？   *  *   *   ？
            # 用*代表墙壁，？代表未翻开的格子，上图左边可以推理出最下面的问号不是地雷，
            # 上图右边可以推理出左边三个问号是地雷，最右边三个问号不是地雷
            for k, v in not_sure.items():
                for i in range(len(v)):
                    for j in range(1, len(v)):
                        s1, n1 = v[i]
                        s2, n2 = v[j]
                        if s1 == s2: continue
                        if s1.issubset(s2):
                            minus_s = s2 - s1
                            if n2 - n1 == minus_s.__len__():
                                for p in minus_s:
                                    to_flag.add(p)
                            elif n2 == n1:
                                for p in minus_s:
                                    to_click.add(p)
                        if s2.issubset(s1):
                            minus_s = s1 - s2
                            if n1 - n2 == minus_s.__len__():
                                for p in minus_s:
                                    to_flag.add(p)
                            elif n2 == n1:
                                for p in minus_s:
                                    to_click.add(p)

                        # 推理模式增强版
                        mixed = s1 & s2
                        if n1 > n2 and s1.__len__() - mixed.__len__() == n1 - n2:
                            for p in s1 - mixed:
                                to_flag.add(p)
                            for p in s2 - mixed:
                                to_click.add(p)
                        if n1 < n2 and s2.__len__() - mixed.__len__() == n2 - n1:
                            for p in s2 - mixed:
                                to_flag.add(p)
                            for p in s1 - mixed:
                                to_click.add(p)

            if action_nums == to_flag.__len__() + to_click.__len__():
                break
    return to_click, to_flag

def game_once(duration=0.01, interval=0.01):
    board = get_cnn_board(model)
    to_click, to_flag = deal_board(board)

    print('click_pos', to_click)
    print('flag_pos', to_flag)
    for lc in sorted(to_click):
        r, c = lc
        pos = (BG_REGION[0] + (c - 0.5) * cols_size, BG_REGION[1] + (r - 0.5) * row_size)
        pyautogui.click(pos, duration=duration)
        pyautogui.sleep(interval)
    for lc in sorted(to_flag):
        r, c = lc
        pos = (BG_REGION[0] + (c - 0.5) * cols_size, BG_REGION[1] + (r - 0.5) * row_size)
        pyautogui.rightClick(pos, duration=duration)
        pyautogui.sleep(interval)
    if not to_click and not to_flag:
        pyautogui.moveTo(1, 1)  # 没有可操作的格子，鼠标提示结束
        return False
    return True

def game_all():
    while game_once():
        pass

def on_press(key):
    # print(key,str(key))
    if isinstance(key, KeyCode):
        if key.char == 'a':
            game_once()

        elif key.char == 'b':
            game_all()

        elif key.char == 'r':
            record_board(model)

        elif key.char == 's':  # 显示识别结果
            pool = Pool(1)
            pool.apply_async(show_recognized, args=(get_cnn_board(model),))
            pool.close()
            pool.join()

        elif key.char == 'c':  # 显示截图结果
            pool = Pool(1)
            pool.apply_async(show_screen_shoot)
            pool.close()
            pool.join()

        elif key.char == 'q':
            return False

def start_listen():
    with Listener(on_press=on_press) as listener:
        listener.join()

if __name__ == '__main__':
    print('开始监听,按 A 键截图并执行脚本一次,\n'
          '按 B 键截图并执行脚本直到无法推理,\n'
          '按 C 键截图并显示截图结果,\n'
          '按 S 键截图并显示预测结果,\n'
          '按 R 键截图并保存截图到文件夹,\n'
          '按 Q 键退出,\n'
          '注：1.如果鼠标失去控制，快速移动鼠标，将鼠标移出屏幕左上角会强制退出程序\n'
          '2.显示截图结果和预测结果与主线程冲突，需要开启新进程，响应会比较慢')
    start_listen()


import importlib
import json
import numpy as np
import os
import sys
import time
import uuid
import copy

from hypergan.discriminators import *
from hypergan.distributions import *
from hypergan.generators import *
from hypergan.inputs import *
from hypergan.samplers import *
from hypergan.trainers import *

import hyperchamber as hc
from hyperchamber import Config
from hypergan.ops import TensorflowOps
import hypergan as hg

from hypergan.gan_component import ValidationException, GANComponent
from ..base_gan import BaseGAN

from hypergan.distributions.uniform_distribution import UniformDistribution

class AlignedAliGAN3(BaseGAN):
    """
    """
    def __init__(self, *args, **kwargs):
        BaseGAN.__init__(self, *args, **kwargs)

    def required(self):
        """
        `input_encoder` is a discriminator.  It encodes X into Z
        `discriminator` is a standard discriminator.  It measures X, reconstruction of X, and G.
        `generator` produces two samples, input_encoder output and a known random distribution.
        """
        return "discriminator ".split()

    def create(self):
        config = self.config
        ops = self.ops

        with tf.device(self.device):
            #x_input = tf.identity(self.inputs.x, name='input')
            xa_input = tf.identity(self.inputs.xa, name='xa_i')
            xb_input = tf.identity(self.inputs.xb, name='xb_i')

            if config.same_g:
                ga = self.create_component(config.generator, input=xb_input, name='a_generator')
                gb = self.create_component(config.generator, input=xa_input, name='a_generator', reuse=True)
            elif config.two_g:
                ga = self.create_component(config.generator1, input=xb_input, name='a_generator')
                gb = self.create_component(config.generator2, input=xa_input, name='b_generator')
            else:
                ga = self.create_component(config.generator, input=xb_input, name='a_generator')
                gb = self.create_component(config.generator, input=xa_input, name='b_generator')

            za = ga.controls["z"]
            zb = gb.controls["z"]

            self.uniform_sample = ga.sample

            xba = ga.sample
            xab = gb.sample
            xa_hat = ga.sample
            xb_hat = gb.sample
            #xa_hat = ga.reuse(gb.sample)
            #xb_hat = gb.reuse(ga.sample)

            z_shape = self.ops.shape(za)
            uz_shape = z_shape
            uz_shape[-1] = uz_shape[-1] // len(config.latent.projections or [1])
            ue = UniformDistribution(self, config.latent, output_shape=uz_shape)
            ue2 = UniformDistribution(self, config.latent, output_shape=uz_shape)
            ue3 = UniformDistribution(self, config.latent, output_shape=uz_shape)
            ue4 = UniformDistribution(self, config.latent, output_shape=uz_shape)
            print('ue', ue.sample)

            zua = ue.sample
            zub = ue2.sample

            uga = ga.sample#ga.reuse(tf.zeros_like(xb_input), replace_controls={"z":zua})
            ugb = gb.sample#gb.reuse(tf.zeros_like(xa_input), replace_controls={"z":zub})

            xa = xa_input
            xb = xb_input

            re_ga = self.create_component(config.generator, input=gb.sample, name='a_generator', reuse=True)
            re_gb = self.create_component(config.generator, input=ga.sample, name='b_generator', reuse=True)
            re_zb = zb#re_gb.controls['z']

            t0 = tf.concat([xb, xb], axis=3)
            t1 = tf.concat([gb.sample, gb.sample], axis=3)
            zaxis = len(self.ops.shape(za))-1

            f0 = tf.concat([za, za], axis=zaxis)
            f1 = tf.concat([zb, zb], axis=zaxis)
            #f0 = tf.concat([zb, za], axis=zaxis)
            #f1 = tf.concat([za, zb], axis=zaxis)
            stack = [t0, t1]
            stacked = ops.concat(stack, axis=0)
            features = ops.concat([f0, f1], axis=0)

            d = self.create_component(config.discriminator, name='d_ab', input=stacked, features=[features])

            self.za = za
            self.discriminator = d
            l = self.create_loss(config.loss, d, None, None, len(stack))
            loss = l
            d1_lambda = config.d1_lambda
            d2_lambda = config.d2_lambda
            d_loss1 = d1_lambda * l.d_loss
            g_loss1 = d1_lambda * l.g_loss

            d_vars1 = d.variables()

            d_loss = l.d_loss
            g_loss = l.g_loss

            metrics = {
                    'g_loss': l.g_loss,
                    'd_loss': l.d_loss
                }

            self._g_vars = ga.variables() + gb.variables()
            self._d_vars = d_vars1
            self.loss=loss
            self.generator = gb
            trainer = self.create_component(config.trainer)
            self.initialize_variables()

        self.trainer = trainer
        self.latent = hc.Config({'sample':zb})
        self.generator = gb
        self.encoder = gb # this is the other gan
        self.uniform_distribution = hc.Config({"sample":zb})#uniform_encoder
        self.zb = zb
        self.z_hat = gb.sample
        self.x_input = xa_input
        self.autoencoded_x = xa_hat

        self.cyca = xa_hat
        self.cycb = xb_hat
        self.xba = xba
        self.xab = xab
        self.uga = uga
        self.ugb = ugb

        rgb = tf.cast((self.generator.sample+1)*127.5, tf.int32)
        self.generator_int = tf.bitwise.bitwise_or(rgb, 0xFF000000, name='generator_int')

    def d_vars(self):
        return self._d_vars
    def g_vars(self):
        return self._g_vars

    def create_discriminator(self, _input, reuse=False):
        return self.create_component(self.config.discriminator, name='d_ab', input=_input, features=[tf.zeros_like(self.za)], reuse=reuse)

    def create_loss(self, loss_config, discriminator, x, generator, split):
        loss = self.create_component(loss_config, discriminator = discriminator, x=x, generator=generator, split=split)
        return loss

    def create_encoder(self, x_input, name='input_encoder'):
        config = self.config
        input_encoder = dict(config.input_encoder or config.g_encoder or config.generator)
        encoder = self.create_component(input_encoder, name=name, input=x_input)
        return encoder

    def create_z_discriminator(self, z, z_hat):
        config = self.config
        z_discriminator = dict(config.z_discriminator or config.discriminator)
        z_discriminator['layer_filter']=None
        net = tf.concat(axis=0, values=[z, z_hat])
        encoder_discriminator = self.create_component(z_discriminator, name='z_discriminator', input=net)
        return encoder_discriminator

    def create_cycloss(self, x_input, x_hat):
        config = self.config
        ops = self.ops
        distance = config.distance or ops.lookup('l1_distance')
        pe_layers = self.gan.skip_connections.get_array("progressive_enhancement")
        cycloss_lambda = config.cycloss_lambda
        if cycloss_lambda is None:
            cycloss_lambda = 10

        if(len(pe_layers) > 0):
            mask = self.progressive_growing_mask(len(pe_layers)//2+1)
            cycloss = tf.reduce_mean(distance(mask*x_input,mask*x_hat))

            cycloss *= mask
        else:
            cycloss = tf.reduce_mean(distance(x_input, x_hat))

        cycloss *= cycloss_lambda
        return cycloss

    def create_z_cycloss(self, z, x_hat, encoder, generator):
        config = self.config
        ops = self.ops
        total = None
        distance = config.distance or ops.lookup('l1_distance')
        if config.z_hat_lambda:
            z_hat_cycloss_lambda = config.z_hat_cycloss_lambda
            recode_z_hat = encoder.reuse(x_hat)
            z_hat_cycloss = tf.reduce_mean(distance(z_hat,recode_z_hat))
            z_hat_cycloss *= z_hat_cycloss_lambda
        if config.z_cycloss_lambda:
            recode_z = encoder.reuse(generator.reuse(z))
            z_cycloss = tf.reduce_mean(distance(z,recode_z))
            z_cycloss_lambda = config.z_cycloss_lambda
            if z_cycloss_lambda is None:
                z_cycloss_lambda = 0
            z_cycloss *= z_cycloss_lambda

        if config.z_hat_lambda and config.z_cycloss_lambda:
            total = z_cycloss + z_hat_cycloss
        elif config.z_cycloss_lambda:
            total = z_cycloss
        elif config.z_hat_lambda:
            total = z_hat_cycloss
        return total

    def input_nodes(self):
        "used in hypergan build"
        if hasattr(self.generator, 'mask_generator'):
            extras = [self.mask_generator.sample]
        else:
            extras = []
        return extras + [
                self.x_input
        ]

    def output_nodes(self):
        "used in hypergan build"

        if hasattr(self.generator, 'mask_generator'):
            extras = [
                self.mask_generator.sample,
                self.generator.g1x,
                self.generator.g2x
            ]
        else:
            extras = []
        return extras + [
                self.encoder.sample,
                self.generator.sample,
                self.uniform_sample,
                self.generator_int
        ]


import torch
import torch.nn.functional as F
import numpy as np

from utils import *
from core.config import config
from train.BaseTrainer import BaseTrainer

class PartNetTrainer(BaseTrainer):

    def __init__(self):
        super(PartNetTrainer, self).__init__()

    def build_opt_and_lr(self, model):

        if config.get('debug'):
            lr_mul = 1
        else:
            lr_mul = len(config.get('gpus'))
        ocfg = config.get('optm_config')
        ignored_params = list(map(id, model.module.resnet.parameters()))
        base_params = filter(lambda p: id(p) not in ignored_params, model.parameters())
        cfg = config.get('dataset_config')
        if cfg['train_name'] == 'market1501':
            new_p_mul = 0.1
        else:
            new_p_mul = 1.
        param_groups = [
            {'params': model.module.resnet.parameters(), 'lr': ocfg['lr'] * new_p_mul * lr_mul},
            {'params': base_params}
                        ]
        if ocfg['name'] == 'SGD':
            optimizer = torch.optim.SGD(param_groups, ocfg['lr'] * lr_mul,
                                        momentum=ocfg['momentum'],
                                        weight_decay=ocfg['weight_decay'])
        else:
            optimizer = torch.optim.Adam(param_groups, ocfg['lr'] * lr_mul,
                                         weight_decay=ocfg['weight_decay'])

        if 'multistep' in ocfg and ocfg['multistep']:
            lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
                                                                ocfg['step'],
                                                                gamma=ocfg['gamma'],
                                                                last_epoch=-1)
        else:
            lr_scheduler = CosineAnnealingWarmUp(optimizer,
                                             T_0=5,
                                             T_end=ocfg.get('epochs'),
                                             warmup_factor=ocfg.get('warmup_factor'),
                                             last_epoch=-1)
        self.logger.write(optimizer)
        return optimizer, lr_scheduler

    def extract(self, test_data, model):
        model.eval()
        res = {}
        for p, val_loader in test_data.items():
            global_features, local_features, local_parts = [], [], []
            with torch.no_grad():
                paths = []
                for i, (input, _, path) in enumerate(val_loader):
                    # print(input[0])
                    input = input.cuda(non_blocking=True)
                    # compute output
                    global_feature, local_feature, local_part = model(input)
                    global_feature = F.normalize(global_feature, p=2, dim=-1)
                    local_feature = F.normalize(local_feature, p=2, dim=-1)

                    if config.get('with_flip'):
                        input_ = input.flip(3)
                        global_feature_, local_feature_, _ = model(input_)
                        global_feature_ = F.normalize(global_feature_, p=2, dim=-1)
                        local_feature_ = F.normalize(local_feature_, p=2, dim=-1)
                        global_feature = (global_feature + global_feature_) / 2
                        local_feature = (local_feature + local_feature_) / 2
                        global_feature = F.normalize(global_feature, p=2, dim=-1)
                        local_feature = F.normalize(local_feature, p=2, dim=-1)

                    global_features.append(global_feature)
                    local_features.append(local_feature)
                    local_parts.append(local_part)

                    paths.extend(path)

                global_features = torch.cat(global_features, dim=0)
                local_features = torch.cat(local_features, dim=0)
                local_parts = torch.cat(local_parts, dim=0)

                print(global_features.size(), local_features.size(), local_parts.size())
                res[p] = {
                    'global_features': global_features,
                    'local_features': local_features,
                    'local_parts': local_parts,
                    'path': paths
                }
        return res

    def eval_status(self, epoch):
        cfg = config.get('dataset_config')
        if cfg['train_name'] == 'market1501':
            return True
        else:
            ocfg = config.get('optm_config')
            return ocfg.get('epochs') - 10 <= epoch <= ocfg.get('epochs')

    def eval_result(self, **kwargs):
        info = kwargs.get('info')
        return evaluate.eval_part_result(info, use_pcb_format=config.get('use_pcb_format', True), logger=self.logger)

    def extract_and_eval(self, test_loader, model):
        res = self.extract(test_loader, model)
        mAP, rank_1 = self.eval_result(info=res)

        return mAP, rank_1

if __name__ == '__main__':
    trainer = PartNetTrainer()
    trainer.train_or_val()


import numpy as np

from elegantrl.demo import *

"""
FinRL PPO
beta1 beta2 gamma reward
beta2 "reward_scaling": 1e-6
beta3 beta2 repeat_times = 2 ** 3

Ant ModSAC Alpha reliable_lambda
ceta2 if if_update_a
ceta3 ceta2 repeat_times = 2 ** 0
"""

def demo4_bullet_mujoco_off_policy():
    args = Arguments(if_on_policy=False)
    args.random_seed = 100867

    from elegantrl.agent import AgentModSAC  # AgentSAC, AgentTD3, AgentDDPG
    args.agent = AgentModSAC()  # AgentSAC(), AgentTD3(), AgentDDPG()
    args.agent.if_use_dn = True
    args.net_dim = 2 ** 7

    import pybullet_envs  # for python-bullet-gym
    dir(pybullet_envs)

    "TotalStep:  3e5, TargetReward: 1500, UsedTime:  8ks, AntBulletEnv-v0"
    "TotalStep:  6e5, TargetReward: 2500, UsedTime: 18ks, AntBulletEnv-v0"
    "TotalStep: 10e5, TargetReward: 2800, UsedTime:   ks, AntBulletEnv-v0"
    args.env = PreprocessEnv(env=gym.make('AntBulletEnv-v0'))
    args.break_step = int(6e5 * 8)  # (5e5) 1e6, UsedTime: (15,000s) 30,000s
    args.if_allow_break = False
    args.reward_scale = 2 ** -2  # todo # RewardRange: -50 < 0 < 2500 < 3340
    args.max_memo = 2 ** 18
    args.batch_size = 2 ** 9
    args.show_gap = 2 ** 8  # for Recorder
    args.eva_size1 = 2 ** 1  # for Recorder
    args.eva_size2 = 2 ** 3  # for Recorder

    # train_and_evaluate(args)
    args.rollout_num = 4
    train_and_evaluate_mp(args)

def check_finrl():
    from FinRL import StockTradingEnv
    from numpy import random as rd
    from finrl.config import config
    from finrl.preprocessing.data import data_split
    import pandas as pd

    # df = pd.read_pickle('finrl_data.df')  # DataFrame of Pandas
    #
    # from finrl.preprocessing.preprocessors import FeatureEngineer
    # fe = FeatureEngineer(
    #     use_technical_indicator=True,
    #     tech_indicator_list=config.TECHNICAL_INDICATORS_LIST,
    #     use_turbulence=True,
    #     user_defined_feature=False,
    # )
    #
    # processed_df = fe.preprocess_data(df)
    # processed_df.to_pickle('finrl_processed_data.df')  # DataFrame of Pandas
    processed_data_path = 'StockTradingEnv_processed_data.df'
    processed_df = pd.read_pickle(processed_data_path)  # DataFrame of Pandas
    print(processed_df.columns.values)

    split_df = data_split(processed_df, start='2008-03-19', end='2021-01-01')
    # `start`

    env = StockTradingEnv(df=split_df, tech_indicator_list=config.TECHNICAL_INDICATORS_LIST)
    action_dim = env.action_dim

    state = env.reset()
    print('state_dim', len(state))

    done = False
    step = 1
    from time import time
    timer = time()
    while not done:
        action = rd.rand(action_dim) * 2 - 1
        next_state, reward, done, _ = env.step(action)
        print(';', step, len(next_state), env.day, reward)
        step += 1

    print(';;', step, int(time() - timer))  # 44 seconds

if __name__ == '__main__':
    # demo2_continuous_action_space_off_policy()
    check_finrl()


import tensorrt as trt
import numpy as np
import pycuda.autoinit
import pycuda.driver as cuda
import time
import sys

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
def build_engine(model_path):
    with trt.Builder(TRT_LOGGER) as builder,builder.create_network(EXPLICIT_BATCH) as network,trt.OnnxParser(network, TRT_LOGGER) as parser:
        builder.max_workspace_size = 1<<20
        builder.max_batch_size = 1
        with open(model_path, "rb") as f:
            parser.parse(f.read())
        engine = builder.build_cuda_engine(network)
    return engine
def inference(engine, context, inputs,h_input, h_output, d_input, d_output,stream):
    cuda.memcpy_htod_async(d_input, h_input, stream)
	  # Run inference.
    context.execute_async(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)
		# Transfer predictions back from the GPU.
    cuda.memcpy_dtoh_async(h_output, d_output, stream)
    return h_output

def alloc_buf(engine):

    h_input = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(0)), dtype=np.float32)
    h_output = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(1)), dtype=np.float32)
    # Allocate device memory for inputs and outputs.
    d_input = cuda.mem_alloc(h_input.nbytes)
    d_output = cuda.mem_alloc(h_output.nbytes)
    stream=cuda.Stream()

    return h_input, h_output, d_input, d_output, stream
if __name__ == "__main__":
    model_path = sys.argv[1] #"/content/gdrive/My Drive/UNIT_Working/onnx-tensorrt-master/result.onnx"
    input_size = int(sys.argv[2])
    output_file=sys.argv[3]
    mode=sys.argv[4]
    if(mode==32):
        inputs = np.random.random((1, 3, input_size, input_size)).astype(np.float32)
        engine = build_engine(model_path)
        print("Engine Created :",type(engine))
        context = engine.create_execution_context()
        print("Context executed ",type(context))
        serialized_engine = engine.serialize()
        t1 = time.time()
        #in_cpu, out_cpu, in_gpu, out_gpu, stream = alloc_buf(engine)
        h_input, h_output, d_input, d_output,stream=alloc_buf(engine)
        res = inference(engine, context, inputs.reshape(-1), h_input, h_output, d_input, d_output, stream)
        #print(type(res))
        #print("using fp32 mode:")
        #print("cost time: ", time.time()-t1)
        #engine_path="FLtask.trt"
        with open(output_file,"wb") as f:
          f.write(serialized_engine)
          print("model converted successfully!")
    if(mode==16):
        inputs = np.random.random((1, 3, input_size, input_size)).astype(np.float16)
        engine = build_engine(model_path)
        print("Engine Created :",type(engine))
        context = engine.create_execution_context()
        print("Context executed ",type(context))
        serialized_engine = engine.serialize()
        t1 = time.time()
        #in_cpu, out_cpu, in_gpu, out_gpu, stream = alloc_buf(engine)
        h_input, h_output, d_input, d_output,stream=alloc_buf(engine)
        res = inference(engine, context, inputs.reshape(-1), h_input, h_output, d_input, d_output, stream)
        #print(type(res))
        #print("using fp16 mode:")
        #print("cost time: ", time.time()-t1)
        with open(output_file,"wb") as f:
          f.write(serialized_engine)
          print("model converted successfully!")
    if(mode==8):
        inputs = np.random.random((1, 3, input_size, input_size)).astype(np.int8)
        engine = build_engine(model_path)
        print("Engine Created :",type(engine))
        context = engine.create_execution_context()
        print("Context executed ",type(context))
        serialized_engine = engine.serialize()
        t1 = time.time()
        #in_cpu, out_cpu, in_gpu, out_gpu, stream = alloc_buf(engine)
        h_input, h_output, d_input, d_output,stream=alloc_buf(engine)
        res = inference(engine, context, inputs.reshape(-1), h_input, h_output, d_input, d_output, stream)
        #print(type(res))
        #print("using int8 mode:")
        #print("cost time: ", time.time()-t1)
        with open(output_file,"wb") as f:
          f.write(serialized_engine)
          print("model converted successfully!")