# Import pygame module
import random
import pygame

# import key coordinates
from pygame.locals import (
    RLEACCEL,
    K_UP,
    K_DOWN,
    K_LEFT,
    K_RIGHT,
    K_ESCAPE,
    KEYDOWN,
    QUIT,
)

# set constant width and height of screen
SCREEN_WIDTH = 800
SCREEN_HEIGHT = 600

# define Player object extending sprite (2D representation of something on screen)
# surface drawn is now attribute of player
class Player(pygame.sprite.Sprite):
    def __init__(self):
        super(Player, self).__init__()
        # Icon Creator: https://www.flaticon.com/authors/dooder
        self.surface = pygame.image.load("my_plane.png").convert()
        self.surface = pygame.transform.scale(self.surface, (60, 60))
        self.surface.set_colorkey((0, 0, 0), RLEACCEL)
        self.rect = self.surface.get_rect()

    def update(self, pressed_keys):
        if pressed_keys[K_UP]:
            self.rect.move_ip(0, -5)
            up_sound.play()
        if pressed_keys[K_DOWN]:
            self.rect.move_ip(0, 5)
            down_sound.play()
        if pressed_keys[K_LEFT]:
            self.rect.move_ip(-5, 0)
        if pressed_keys[K_RIGHT]:
            self.rect.move_ip(5, 0)

        # keep player on screen
        if self.rect.left < 0:
            self.rect.left = 0
        if self.rect.right > SCREEN_WIDTH:
            self.rect.right = SCREEN_WIDTH
        if self.rect.top <= 0:
            self.rect.top = 0
        if self.rect.bottom >= SCREEN_HEIGHT:
            self.rect.bottom = SCREEN_HEIGHT

class Enemy(pygame.sprite.Sprite):
    def __init__(self):
        super(Enemy, self).__init__()
        # Icon Creator: https://www.flaticon.com/authors/flat-icons
        self.surface = pygame.image.load("enemy.png").convert()
        self.surface = pygame.transform.scale(self.surface, (20, 20))
        self.surface.set_colorkey((0, 0, 0), RLEACCEL)
        self.rect = self.surface.get_rect(
            center=(
                random.randint(SCREEN_WIDTH + 20, SCREEN_WIDTH + 100),
                random.randint(0, SCREEN_HEIGHT),
            )
        )
        self.speed = random.randint(5, 20)

    # move sprite based on speed
    # remove sprite after passing left edge of screen
    def update(self):
        self.rect.move_ip(-self.speed, 0)
        if self.rect.right < 0:
            self.kill()

# cloud object with image
class Cloud(pygame.sprite.Sprite):
    def __init__(self):
        super(Cloud, self).__init__()
        self.surface = pygame.image.load("cloud.png").convert()
        self.surface.set_colorkey((0, 0, 0), RLEACCEL)
        # start pos random
        self.rect = self.surface.get_rect(
            center=(
                random.randint(SCREEN_WIDTH + 20, SCREEN_WIDTH + 100),
                random.randint(0, SCREEN_HEIGHT),
            )
        )

    # move cloud on constant speed, remove after passing left edge of screen
    def update(self):
        self.rect.move_ip(-5, 0)
        if self.rect.right < 0:
            self.kill()

# setup sounds
pygame.mixer.init()

# Initialize pygame
pygame.init()

# setup clock for decent framerate
clock = pygame.time.Clock()

# create screen object of size SCREEN_WIDTH and SCREEN_HEIGHT
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))

# create custom event to add new enemy
ADDENEMY = pygame.USEREVENT + 1
pygame.time.set_timer(ADDENEMY, 250)
ADDCLOUD = pygame.USEREVENT + 2
pygame.time.set_timer(ADDCLOUD, 1000)

# instantiate player
my_player = Player()

# create sprite group, an object holding group of Sprite objects to hold enemy sprites
# used for collision detection and position updates
# all_sprites to render
enemies = pygame.sprite.Group()
clouds = pygame.sprite.Group()
all_sprites = pygame.sprite.Group()
all_sprites.add(my_player)

# background music
# Sound source: http://ccmixter.org/files/Apoxode/59262
# License: https://creativecommons.org/licenses/by/3.0/
pygame.mixer.music.load("Apoxode_-_Electric_1.mp3")
pygame.mixer.music.play(loops=-1)

# sound files
# sound source: Jon Fincher
up_sound = pygame.mixer.Sound("Rising_putter.ogg")
down_sound = pygame.mixer.Sound("Falling_putter.ogg")
collision_sound = pygame.mixer.Sound("Collision.ogg")

# Set base volume
up_sound.set_volume(0.5)
down_sound.set_volume(0.5)
collision_sound.set_volume(0.5)

# keep main loop running aka variable to keep game running
running = True

# Main loop
while running:
    # look at every event in queue
    for event in pygame.event.get():
        # if user hits a key
        if event.type == KEYDOWN:
            # check if escape key to stop loop
            if event.key == K_ESCAPE:
                running = False

        # if user closes window, stop loop also
        elif event.type == QUIT:
            running = False

        # add new enemy
        elif event.type == ADDENEMY:
            # create new enemy, add to sprite group
            new_enemy = Enemy()
            enemies.add(new_enemy)
            all_sprites.add(new_enemy)

        # add new cloud
        elif event.type == ADDCLOUD:
            # create new cloud and add to sprite group
            new_cloud = Cloud()
            clouds.add(new_cloud)
            all_sprites.add(new_cloud)

    # get keys pressed, check for user input
    pressed_keys = pygame.key.get_pressed()

    # update my player sprite based on user key presses
    my_player.update(pressed_keys)

    # update enemy position and cloud pos every frame
    enemies.update()
    clouds.update()

    # make screen white
    screen.fill((135, 206, 250))

    # draw all sprites
    for charSprite in all_sprites:
        screen.blit(charSprite.surface, charSprite.rect)

    # check if any enemies collided with player
    if pygame.sprite.spritecollideany(my_player, enemies):
        # remove player and stop loop -> game ends
        my_player.kill()

        # stop moving sounds, play collision sound
        up_sound.stop()
        down_sound.stop()
        collision_sound.play()

        running = False

    # update display
    pygame.display.flip()

    # ensure program has rate of 30 frames per second
    clock.tick(30)

# stop and quit mixer
pygame.mixer.music.stop()
pygame.mixer.quit()


import paho.mqtt.client as mqttclient
import time
import random
from pymongo import MongoClient
import os

import os
# from dotenv import load_dotenv
# load_dotenv()
connection_string = os.environ.get("MONGO_CONNECTION_STRING") # environment vars

# 50 per frame

def on_connect(client, userdata, flags, rc):
    if rc==0:
        print("client is connected")
        global connected
        connected=True
    else:
        print("client is not connected")

def on_message(client, userdata, message):

    try:
        read_message = message.payload.decode("utf-8")
        if read_message[0] == '[':
            frames = read_message[1:-1].split(',')
        else:
            frames = read_message.split(',')
        # print(frames)
        data_to_DB = []
        for index, frame in enumerate(frames):
            arr = frame.split()
            a = {'device': arr[0], 'wind_speed':arr[3], '_id': str(arr[1]) + " " + str(arr[2]),
            'wind_heading': arr[4], 'pm1':arr[5], 'pm25':arr[6], 'pm10':arr[7]}
            data_to_DB.append(a)
            # print(str(arr[1]) + " " + str(arr[2])) #this is the "_id" field in mongo doc

        print(data_to_DB)
        collection.insert_many(data_to_DB)
        print("Topic: " + str(message.topic))

    except:
        pass

# Set up client for MongoDB
# mongoClient=MongoClient("mongodb://localhost:27017/")
# connection_string = os.getenv('MONGO_CONNECTION_STRING')

mongoClient=MongoClient(connection_string)
db=mongoClient.SensorData
collection=db.sensor1

connected = False
MessageReceived = False

broker_address = 'broker.hivemq.com'

port = 1883
client_id = f'python-mqtt-{random.randint(0, 100)}'

client = mqttclient.Client(client_id)
client.on_connect = on_connect
client.connect(broker_address, port)
client.loop_start()

client.subscribe("praan/mqtt")
while True:
    # time.sleep(1)
    client.on_message = on_message

# client.loop_end()


import pygame
pygame.init()

Font = pygame.font.Font('Pixeltype.ttf', 50)

def BUTTON(content, color, pos, size):
    rect = pygame.Rect(pos[0] - size[0]//2 , pos[1] - size[1]//2, size[0], size[1])
    text = Font.render(content, False, color)
    return text, rect

class Menu:
    def __init__(self, game):
        self.game = game
        self.server, self.ser_rect = BUTTON("sever", 'orange', (320,200), (110, 50))
        self.client, self.cl_rect = BUTTON('Client', "Blue", (320, 260), (110, 50))

    def run(self, surf):
        self.render(surf)

        mouse = pygame.mouse.get_pos()

        if self.ser_rect.collidepoint(mouse[0], mouse[1]):
            pygame.mouse.set_system_cursor(pygame.SYSTEM_CURSOR_HAND)
        elif self.cl_rect.collidepoint(mouse[0], mouse[1]):
            pygame.mouse.set_system_cursor(pygame.SYSTEM_CURSOR_HAND)
        else:
             pygame.mouse.set_system_cursor(0)

        mouse_click = pygame.mouse.get_pressed()
        if mouse_click[0]:
            if self.ser_rect.collidepoint(mouse[0], mouse[1]):
                pygame.mouse.set_system_cursor(0)
                return 'server'
            if self.cl_rect.collidepoint(mouse[0], mouse[1]):
                pygame.mouse.set_system_cursor(0)
                return 'client'

    def render(self, surf):
        ## SERVER
        sur_surf = pygame.Surface((self.ser_rect.width, self.ser_rect.height))
        sur_surf.fill('black')
        surf.blit(sur_surf, self.ser_rect)
        text_rect = self.server.get_rect(center = self.ser_rect.center)
        surf.blit(self.server, text_rect)

        ##CLIENT
        sur_surf = pygame.Surface((self.cl_rect.width, self.cl_rect.height))
        sur_surf.fill('black')
        surf.blit(sur_surf, self.cl_rect)
        text_rect = self.client.get_rect(center = self.cl_rect.center)
        surf.blit(self.client, text_rect)


import os
import numpy as np
import shutil
import random
import math
import pathlib

import sys
curr_dir = os.path.dirname(os.path.abspath(__file__))
root_dir = os.path.join(curr_dir, '..')
external_dir = os.path.join(root_dir, 'externals')
sys.path.insert(0, root_dir)
sys.path.insert(1, os.path.join(external_dir, 'pytorch_a2c_ppo_acktr_gail'))

from ppo import run_ppo
from evogym import sample_robot, hashable
import utils.mp_group as mp
from utils.algo_utils import get_percent_survival_evals, mutate, Structure
from NestedOptimization import NestedOptimization

from ppo.envs import make_vec_envs
import torch
from ppo.utils import get_vec_normalize
import imageio
from pygifsicle import optimize
import pickle

def dump_visualization_data(dump_path, out_gif_path, env_name, structure, ctrl_path):

    out_gif_path, env_name, structure, ctrl_path

    print("Saving simulation objects...", end="")

    # an example object to be pickled
    simulation_objects = [out_gif_path, env_name, structure, ctrl_path]
    print("ctrl_path =",ctrl_path)
    # pickling the object
    with open(dump_path, "wb") as f:
        pickle.dump(simulation_objects, f)
    print("done.")

def load_visualization_data(pickle_dump_path):
    with open(pickle_dump_path, "rb") as f:
        out_path, env_name, structure, ctrl_path = pickle.load(f)

    return out_path, env_name, structure, ctrl_path

def save_robot_gif_standalone(out_path, env_name, structure, ctrl_path):
    print("Generating gif...")
    gif_resolution = (1280/5, 720/5)

    env = make_vec_envs(env_name, structure, 1000, 1, None, None, device='cpu', allow_early_resets=False)
    env.get_attr("default_viewer", indices=None)[0].set_resolution(gif_resolution)

    actor_critic, obs_rms = torch.load(ctrl_path, map_location='cpu')

    vec_norm = get_vec_normalize(env)
    if vec_norm is not None:
        vec_norm.eval()
        vec_norm.obs_rms = obs_rms

    recurrent_hidden_states = torch.zeros(1, actor_critic.recurrent_hidden_state_size)
    masks = torch.zeros(1, 1)

    obs = env.reset()
    img = env.render(mode='img')
    reward = None
    done = False

    imgs = []
    # arrays = []
    while not done:

        with torch.no_grad():
            value, action, _, recurrent_hidden_states = actor_critic.act(
                obs, recurrent_hidden_states, masks, deterministic=True)

        obs, reward, done, _ = env.step(action)
        img = env.render(mode='img')
        imgs.append(img)

        masks.fill_(0.0 if (done) else 1.0)

        if done == True:
            env.reset()

    env.close()
    imageio.mimsave(out_path, imgs, duration=(1/50.0))

def run_ga(pop_size, structure_shape, no: NestedOptimization):

    seed = no.params.seed
    env_name = no.params.env_name
    max_evaluations = no.params.max_frames

    random.seed(seed)
    import numpy as np
    np.random.seed(seed)

    try:
        os.makedirs( os.path.join(root_dir, "saved_data"))
        os.makedirs( os.path.join(root_dir, "saved_data", f"exp{no.params.experiment_index}"))
    except:
        pass

    internal_exp_files = os.path.join(root_dir, "saved_data", f"exp{no.params.experiment_index}")
    temp_path = internal_exp_files + "/metadata.txt"

    if os.path.isdir(internal_exp_files):
        print("Removing old exp. files:")
        import shutil
        shutil.rmtree(internal_exp_files, ignore_errors=True)

    os.makedirs(internal_exp_files)

    f = open(temp_path, "w")
    f.write(f'POP_SIZE: {pop_size}\n')
    f.write(f'STRUCTURE_SHAPE: {structure_shape[0]} {structure_shape[1]}\n')
    f.write(f'MAX_EVALUATIONS: {max_evaluations}\n')
    f.close()

    ### GENERATE // GET INITIAL POPULATION ###
    structures = []
    population_structure_hashes = {}
    num_evaluations = 0
    generation = 0

    #generate a population
    for i in range (pop_size):

        temp_structure = sample_robot(structure_shape)
        while (hashable(temp_structure[0]) in population_structure_hashes):
            temp_structure = sample_robot(structure_shape)

        structures.append(Structure(*temp_structure, i))
        population_structure_hashes[hashable(temp_structure[0])] = True
        num_evaluations += 1

    while True:

        ### UPDATE NUM SURVIORS ###
        # We need to update the percentage survival based on the frames left and not on number of evaluations left.
        percent_survival = get_percent_survival_evals(no.step, no.max_frames)
        num_survivors = max(2, math.ceil(pop_size * percent_survival))

        ### MAKE GENERATION DIRECTORIES ###
        save_path_structure = os.path.join(root_dir, "saved_data", f"exp{no.params.experiment_index}", "generation_" + str(generation), "structure")
        save_path_controller = os.path.join(root_dir, "saved_data", f"exp{no.params.experiment_index}", "generation_" + str(generation), "controller")

        try:
            os.makedirs(save_path_structure)
        except:
            pass

        try:
            os.makedirs(save_path_controller)
        except:
            pass

        ### SAVE POPULATION DATA ###
        for i in range (len(structures)):
            temp_path = os.path.join(save_path_structure, str(structures[i].label))
            np.savez(temp_path, structures[i].body, structures[i].connections)

        for structure in structures:

            if structure.is_survivor:
                save_path_controller_part = os.path.join(root_dir, "saved_data", f"exp{no.params.experiment_index}", "generation_" + str(generation), "controller",
                    "robot_" + str(structure.label) + "_controller" + ".pt")
                print(f'Skipping training for {save_path_controller_part}.\n')
            else:

                controller_path_for_animation_current = f"controller_to_generate_animation_{no.params.experiment_index}_current.pt"
                controller_path_for_animation_best = f"controller_to_generate_animation_{no.params.experiment_index}_best.pt"
                no.controller_path_for_animation = controller_path_for_animation_current

                # For sequential execution
                res = run_ppo((structure.body, structure.connections), (save_path_controller, structure.label), env_name, no, False)
                structure.set_reward(res)

                if no.is_reevaluating_flag:
                    _ = run_ppo((structure.body, structure.connections), (save_path_controller, structure.label), env_name, no, True)  # no.new_best_found is updated here
                    import pathlib
                    dump_path_current = f"simulation_objects_{no.params.experiment_index}_current.pkl"
                    out_path_gif_current = pathlib.Path().resolve().as_posix() + f"../../../../results/evogym/videos/vid_{no.get_video_label()}_current.gif"
                    dump_visualization_data(dump_path_current, out_path_gif_current, env_name, (structure.body, structure.connections), controller_path_for_animation_current)

                    if no.new_best_found:
                        dump_path_best = f"simulation_objects_{no.params.experiment_index}_best.pkl"
                        out_path_gif_best = pathlib.Path().resolve().as_posix() + f"../../../../results/evogym/videos/vid_{no.get_video_label()}_best.gif"
                        os.system(f"cp {controller_path_for_animation_current} {controller_path_for_animation_best}")
                        dump_visualization_data(dump_path_best, out_path_gif_best, env_name, (structure.body, structure.connections), controller_path_for_animation_best)
                        no.new_best_found = False

        ### COMPUTE FITNESS, SORT, AND SAVE ###
        for structure in structures:
            structure.compute_fitness()

        structures = sorted(structures, key=lambda structure: structure.fitness, reverse=True)

        #SAVE RANKING TO FILE
        temp_path = os.path.join(root_dir, "saved_data", f"exp{no.params.experiment_index}", "generation_" + str(generation), "output.txt")
        f = open(temp_path, "w")

        out = ""
        for structure in structures:
            out += str(structure.label) + "\t\t" + str(structure.fitness) + "\n"
        f.write(out)
        f.close()

        # print(f'FINISHED GENERATION {generation} - SEE TOP {round(percent_survival*100)} percent of DESIGNS:\n')
        # print(structures[:num_survivors])

        ### CROSSOVER AND MUTATION ###
        # save the survivors
        survivors = structures[:num_survivors]

        #store survivior information to prevent retraining robots
        for i in range(num_survivors):
            structures[i].is_survivor = True
            structures[i].prev_gen_label = structures[i].label
            structures[i].label = i

        # for randomly selected survivors, produce children (w mutations)
        num_children = 0
        while num_children < (pop_size - num_survivors) and num_evaluations < max_evaluations:

            parent_index = random.sample(range(num_survivors), 1)
            child = mutate(survivors[parent_index[0]].body.copy(), mutation_rate = 0.1, num_attempts=50)

            if child != None and hashable(child[0]) not in population_structure_hashes:

                # overwrite structures array w new child
                structures[num_survivors + num_children] = Structure(*child, num_survivors + num_children)
                population_structure_hashes[hashable(child[0])] = True
                num_children += 1
                num_evaluations += 1

        structures = structures[:num_children+num_survivors]

        generation += 1


from __future__ import print_function
import os,sys
import numpy as np
import math
import siconos.kernel as Kernel
np.set_printoptions(precision=3)

from siconos.mechanics.collision.tools import Contactor
from siconos.mechanics.joints import cast_PivotJointR
from siconos.io.mechanics_run import MechanicsHdf5Runner
from siconos.kernel import SiconosVector, BlockVector, changeFrameAbsToBody

import siconos.numerics as sn
import siconos.kernel as sk

# An example of applying force to the axis of a joint, and applying
# spring and virtual damping by measuring position and velocity along
# the same axis.

# Note: This example is to demonstrate external measurement of joint
# positions and application of forces to dynamical systems attached to
# joints.  In practice it is better to use internal forces (fInt,
# mInt) to model joint spring-dampers, see folder
# JointsTestsWithInternalForces, and extra Relations with associated
# Non-Smooth Laws to model non-linearities such as joint stops and
# friction, see JointsTestsWithContactDetection.

# Creation of the hdf5 file for input/output
with MechanicsHdf5Runner() as io:

    # Definition of two bars connected by a prismatic joint
    io.add_primitive_shape('Bar', 'Box', (1, 0.1, 0.1))
    io.add_object('bar1', [Contactor('Bar')], [0,0,2],
                 orientation=[(0,0,1),np.pi/2], mass=1.0, velocity=[0,0,0,0,0,1])
    io.add_object('bar2', [Contactor('Bar',relative_translation=[0, 0.1,0]),
                          Contactor('Bar',relative_translation=[0,-0.1,0])],
                 [0,0,2],
                 orientation=[(0,0,1),np.pi/2], mass=1.0)
    io.add_joint('joint1', 'bar1', 'bar2', [[0,0,0]], [[0,1,0]], 'PivotJointR',
                absolute=False)

    # Definition of the ground
    io.add_primitive_shape('Ground', 'Box', (5, 5, 0.1))
    io.add_object('ground', [Contactor('Ground')], [0,0,-0.05])
    io.add_Newton_impact_friction_nsl('contact', mu=0.3, e=0.0)

class Ctrl(object):
    def initialize(self, io):
        self.count = 0
        self.topo = io._nsds.topology()
        self.ds1 = Kernel.cast_NewtonEulerDS(self.topo.getDynamicalSystem('bar1'))
        self.ds2 = Kernel.cast_NewtonEulerDS(self.topo.getDynamicalSystem('bar2'))
        self.joint1 = cast_PivotJointR(
            self.topo.getInteraction('joint1').relation())

        self.ds1.setIsMextExpressedInInertialFrame(True)
        self.ds2.setIsMextExpressedInInertialFrame(True)

        # Apply initial forces
        self.step()

    def step(self):
        self.count += 1

        # Make a temporary BlockVector containing both qs
        bv = BlockVector(self.ds1.q(), self.ds2.q())

        torque1 = np.zeros(3)
        torque2 = np.zeros(3)

        # Get the position and use it to project a torque vector
        # onto the DoF (spring torque)
        angle = SiconosVector(1)
        self.joint1.computehDoF(0, bv, angle, 0)

        setpoint = np.pi/4
        ang_diff = setpoint - angle.getValue(0)
        spring_torque = np.array(self.joint1.normalDoF(bv, 0)) * ang_diff * 500.0

        # Get the velocity of each body projected onto the DoF and
        # calculate their difference (damping torque)
        vel1 = self.joint1.projectVectorDoF(self.ds1.angularVelocity(True), bv, 0)
        vel2 = self.joint1.projectVectorDoF(self.ds2.angularVelocity(True), bv, 0)
        vel_diff = vel1 - vel2
        damping_torque = vel_diff * 5.0

        # Calculate total torques for each body
        torque1 += - (spring_torque + damping_torque)/2
        torque2 += + (spring_torque + damping_torque)/2

        print('applying spring-damper torques', torque1, torque2)

        self.ds1.setMExtPtr(torque1)
        self.ds2.setMExtPtr(torque2)

options = sk.solver_options_create(sn.SICONOS_GENERIC_MECHANICAL_NSGS)
options.iparam[sn.SICONOS_IPARAM_MAX_ITER] = 1000
options.dparam[sn.SICONOS_DPARAM_TOL] = 1e-12

# Load and run the simulation
with MechanicsHdf5Runner(mode='r+') as io:
    io.run(t0=0,
           T=5,
           h=0.001,
           theta=0.5,
           Newton_max_iter=1,
           controller=Ctrl(),
           solver_options=options,
           output_frequency=1)


# Script to read in neuron data and do correlations
# First do correlation for all neuronxall neuron
# Using that matrix of data, pull out the data of AVA specific correlations

# For data set 110803 (moving only)- frames 1-1465, AVA 33 and 16

# For data set 102246 (moving-immobile)-Moving 1-1590, Immobile 1700-4016 (maybe less/still immobile then), AVA 29 (30 matlab) (and 17?)
# For data set 111126 (moving-immobile)-moving 1-1785, Immobile 1900-3245, AVA 27 (28 matlab)
# For data set 144610 (moving-immobile)-moving 0-1429, Immobile 1539-2904, BFP 2924-3nd, AVA- 17 and 81

#For Dataset 173358-immobile only curre

import os
import numpy as np
import matplotlib.pyplot as plt
from prediction import userTracker
import prediction.dataHandler as dh
from seaborn import clustermap

AVA_neuron1 = 72 #for dataset 144610
AVA_neuron2 = 22 #for dataset 144610

def get_data(start_frame, end_frame, neurons, neurons_withNaN, neurons_Zscore, time, time_contig, velocity):
    # take the loaded in data and pull out the neuron data for moving or immobile.
    # look up time of start frame/end frame in the time_contig
    # see what index in time (noncontig) is closest to time_contig value
    start_time= time_contig[start_frame]
    start_index = np.abs(time - start_time).argmin()
    end_time = time_contig[end_frame]
    end_index = np.abs(time-end_time).argmin()
    neuron_data = neurons[:, start_index:end_index]
    neuron_NaN_data = neurons_withNaN[:, start_index:end_index]
    zscore_neuron = neurons_Zscore[:, start_index:end_index]
    [row, column] = neurons.shape
    neuron_num = list(range(0, row, 1))
    crop_velocity = velocity[start_frame:end_frame]
    crop_time = time_contig[start_frame:end_frame]
    return neuron_data, neuron_NaN_data, zscore_neuron, neuron_num, crop_velocity, crop_time

def do_correlation_all(neuron_data):
    # take in the Ratio2 frames, do the correlations/make the matrix
    r = np.corrcoef(neuron_data)
    # square the matrix to get the r_square
    r_square = np.square(r)
    return r, r_square

def find_NaN_neurons(neuron_NaN_data, neuron_num):
    # step through each row of data and find the sum of NaNs. If majority NaN, save that index (neuron number)
    index_list = []
    for i in neuron_num:
        row = neuron_NaN_data[i,:]  # just pull out the row we want, then check if it is majority NaNs
        column = len(row)
        numNan = np.sum(np.isnan(row))
        ratio = float(numNan)/float(column)
        if ratio > 0.8:  # for the non-tracked neurons, all time points should be NaN, so ratio should be 1
            index_list.append(i)
    return index_list

def replace_NaN_neurons(r, index_list, neuron_num):
    # for the moving-immobile transition, use index list to see where non tracked neurons are. Replace those
    # rows/columns with 0s in the r/r_square
    for i in index_list:
        # each index in the list is a row of NaNs. replace the correlation with 0s if neuron isn't tracked
        r[i, :] = np.zeros(len(neuron_num))
        r[:, i] = np.zeros(len(neuron_num))
    return r

def do_correlation_AVA(r, r_square, AVA_neuron):
    # take the big all correlations and find the AVA neuron
    AVA_correlation = r[AVA_neuron, :]
    AVA_r_square = r_square[AVA_neuron,:]
    return AVA_correlation, AVA_r_square

def replace_Nan_AVA(AVA_list, index_list):
    for i in index_list:
        AVA_list[i] = 0
    return AVA_list

def insert_Nan_neurons(neuron_list, index_list):
    # replace the untracked neurons immobile data with NaN in the moving data
    for i in index_list:
        neuron_list[i, :] = np.NaN
    return neuron_list

def do_residual(moving_worm, immobile_worm):
    # this function should take the residual of whatever we give it (should work for whole matrix correlation as well
    # as AVA specific correlation)
    residual = immobile_worm-moving_worm
    return residual

def do_clustering(matrix):
    cg = clustermap(matrix)
    cgIdx = np.array(cg.data2d.columns)
    #cgIdx is a list of the resorted neuron numbers. Want to mark neurons for AVA
    return cgIdx

def get_index_AVA_clustered(AVAneuron, cgIdx):
    #take the cgIdx and find what index corresponds to the AVA neuron
    AVA_index = np.where(cgIdx == AVAneuron)
    return AVA_index

def cluster_matrix(matrix, cgIdx):
    clustered_1 = matrix[:, cgIdx]
    clustered_matrix = clustered_1[cgIdx, :]
    return clustered_matrix

def remove_nan(matrix, index_list):
    # remove the matrix rows/columns that are the Nan values
    # remove the nan_rows
    matrix_no_row_nan = np.delete(matrix, index_list,0)
    # remove the nan_columns
    matrix_no_nan = np.delete(matrix_no_row_nan, index_list, 1)
    return matrix_no_nan

def get_index_after_NaN(old_index, index_list):
    #Take the full dataset and count how many NaNs will be removed to get to index of interest (AVA). Subtract accordingly
    count = len([i for i in index_list if i < old_index])
    new_index = old_index - count
    return new_index

def remove_and_cluster_and_residual(matrix_move, matrix_immobile, index_list):
    matrix_move_noNaN = remove_nan(matrix_move, index_list)
    matrix_immobile_noNaN = remove_nan(matrix_immobile, index_list)
    #do immobile clustering
    immobile_cgIdx = do_clustering(matrix_immobile_noNaN)
    moving_cgIdx = do_clustering(matrix_move_noNaN)
    new_AVA_index_noNaN1 = get_index_after_NaN(AVA_neuron1, index_list)
    new_AVA_index_noNaN2 = get_index_after_NaN(AVA_neuron2, index_list)
    movecluster_AVA_index1 = get_index_AVA_clustered(new_AVA_index_noNaN1, moving_cgIdx)
    movecluster_AVA_index2 = get_index_AVA_clustered(new_AVA_index_noNaN2,moving_cgIdx)
    movecluster_AVA = [movecluster_AVA_index1, movecluster_AVA_index2]
    immcluster_AVA_index1 = get_index_AVA_clustered(new_AVA_index_noNaN1, immobile_cgIdx)
    immcluster_AVA_index2 = get_index_AVA_clustered(new_AVA_index_noNaN2,immobile_cgIdx)
    immcluster_AVA = [immcluster_AVA_index1, immcluster_AVA_index2]
    movecluster_move_noNaN = cluster_matrix(matrix_move_noNaN, moving_cgIdx)
    movecluster_immobile_noNaN = cluster_matrix(matrix_immobile_noNaN, moving_cgIdx)
    immcluster_move_noNaN = cluster_matrix(matrix_move_noNaN, immobile_cgIdx)
    immcluster_immobile_noNaN = cluster_matrix(matrix_immobile_noNaN, immobile_cgIdx)
    movecluster_residual = do_residual(movecluster_move_noNaN, movecluster_immobile_noNaN)
    immcluster_residual = do_residual(immcluster_move_noNaN, immcluster_immobile_noNaN)
    return movecluster_move_noNaN, movecluster_immobile_noNaN, movecluster_residual, movecluster_AVA, immcluster_move_noNaN, immcluster_immobile_noNaN, immcluster_residual, immcluster_AVA

def sum_matrix(matrix):
    # sum together all of the elements of the matrix. subtract the 1:1 correlation, divide by matrix
    #take abs value of matrix
    abs_matrix = np.abs(matrix)
    total = abs_matrix.sum()
    diagonal = len(matrix)
    corr = total-diagonal # subtract off the 1:1
    elements = (diagonal-1)*diagonal
    avg_corr = corr/elements
    return avg_corr

# Want to cluster the calcium signal- use correlation clustering?
# sort the NaN data?
def cluster_calcium(neurons_NaN, cgIdx):
    cluster_neuron = neurons_NaN[cgIdx, :]
    return cluster_neuron

for typ_cond in ['AML310_transition']: #, 'AML310_moving']:
    path = userTracker.dataPath()
    folder = os.path.join(path, '%s/' % typ_cond)
    dataLog = os.path.join(path,'{0}/{0}_datasets.txt'.format(typ_cond))
    # data parameters
    dataPars = {'medianWindow': 0,  # smooth eigenworms with gauss filter of that size, must be odd
            'gaussWindow': 50,  # gaussianfilter1D is uesed to calculate theta dot from theta in transformEigenworms
            'rotate': False,  # rotate Eigenworms using previously calculated rotation matrix
            'windowGCamp': 5,  # gauss window for red and green channel
            'interpolateNans': 6,  # interpolate gaps smaller than this of nan values in calcium data
            'volumeAcquisitionRate': 6.,  # rate at which volumes are acquired
            }
    dataSets = dh.loadMultipleDatasets(dataLog, pathTemplate = folder, dataPars = dataPars)
    keyList = np.sort(dataSets.keys())
    for key in filter(lambda x: '193044' in x, keyList):
        print("Running "+key)
        time = dataSets[key]['Neurons']['I_Time_crop_noncontig']
        time_contig = dataSets[key]['Neurons']['I_Time']
        neurons = dataSets[key]['Neurons']['I_smooth_interp_crop_noncontig']
        neurons_withNaN = dataSets[key]['Neurons']['I_smooth'] # use this to find the untracked neurons after transition
        neurons_ZScore = dataSets[key]['Neurons']['ActivityFull'] # Z scored neurons to use to look at calcium traces
        velocity = dataSets[key]['Behavior_crop_noncontig']['AngleVelocity']
        # curvature = dataSets[key]['Behavior_crop_noncontig']['Eigenworm3']
        # For immobile- how is NaN neurons that are not hand tracked dealt with by the smooth_interp...
        # Still do the correlation with all (the interpolated ones too, but then replace with 0s)?

    transition = im_start = 950
    im_end = 2885

    moving_data, moving_withNaN, moving_zscore, neuron_number, moving_velocity, moving_time = get_data(0, transition, neurons, neurons_withNaN, neurons_ZScore, time, time_contig,velocity)
    moving_corr, moving_r_square = do_correlation_all(moving_data)
    AVA_corr_noNan1, AVA_r_square_noNan1 = do_correlation_AVA(moving_corr,moving_r_square, AVA_neuron1)
    AVA_corr_noNan2, AVA_r_square_noNan2 = do_correlation_AVA(moving_corr,moving_r_square, AVA_neuron2)

    # cgIdx = do_clustering(moving_corr_NaN) cluster based on immobile

    immobile_data, data_withNaN, immobile_zscore, neuron_number_im, imm_velocity, imm_time = get_data(transition, im_end, neurons, neurons_withNaN, neurons_ZScore, time, time_contig, velocity)
    immobile_corr, immobile_r_square = do_correlation_all(immobile_data)
    AVA_corr_im1, AVA_r_square_im1 = do_correlation_AVA(immobile_corr, immobile_r_square, AVA_neuron1)
    AVA_corr_im2, AVA_r_square_im2 = do_correlation_AVA(immobile_corr, immobile_r_square, AVA_neuron2)
    nan_index = find_NaN_neurons(data_withNaN, neuron_number)
    immobile_corr_NaN = replace_NaN_neurons(immobile_corr, nan_index, neuron_number) # Nan = 0
    immobile_r_NaN = replace_NaN_neurons(immobile_r_square, nan_index, neuron_number)

    moving_corr_NaN = replace_NaN_neurons(moving_corr, nan_index, neuron_number)
    moving_r_square_NaN = replace_NaN_neurons(moving_r_square, nan_index, neuron_number)

    #Remove Nan, cluster and do residuals on correlation matrix
    mcluster_move_noNaN, mcluster_imm_noNaN, mcluster_residual, mcluster_AVA, imcluster_move_noNaN, imcluster_imm_noNaN, imcluster_residual, imcluster_AVA = remove_and_cluster_and_residual(moving_corr, immobile_corr, nan_index)
    move_noNaN_avg_corr = sum_matrix(mcluster_move_noNaN)
    imm_nonNaN_avg_corr = sum_matrix(mcluster_imm_noNaN)
    move_noNaN_avg_corr = round(move_noNaN_avg_corr,2)
    imm_nonNaN_avg_corr = round(imm_nonNaN_avg_corr,2)

    #Put NaN/0s into the AVA measurements as well, propagating them to the moving part
    AVA_corr1 = replace_Nan_AVA(AVA_corr_noNan1, nan_index)
    AVA_r_square1 = replace_Nan_AVA(AVA_r_square_noNan1, nan_index)

    AVA_corr2 = replace_Nan_AVA(AVA_corr_noNan2, nan_index)
    AVA_r_square2 = replace_Nan_AVA(AVA_r_square_noNan2, nan_index)

    cgIdx_immobile = do_clustering(immobile_corr_NaN)
    clustered_immobile_corr = cluster_matrix(immobile_corr_NaN, cgIdx_immobile)
    clustered_immobile_r = cluster_matrix(immobile_r_NaN, cgIdx_immobile)
    AVA_index_imm1 = get_index_AVA_clustered(AVA_neuron1, cgIdx_immobile)
    AVA_index_imm2 = get_index_AVA_clustered(AVA_neuron2, cgIdx_immobile)

    clustered_moving_corr = cluster_matrix(moving_corr_NaN, cgIdx_immobile)
    clustered_moving_r = cluster_matrix(moving_r_square_NaN, cgIdx_immobile)

    # cluster the neuron signal
    data_withNaN = insert_Nan_neurons(data_withNaN, nan_index)
    moving_withNaN = insert_Nan_neurons(moving_withNaN, nan_index)
    immobile_zscore = insert_Nan_neurons(immobile_zscore, nan_index)
    moving_zscore = insert_Nan_neurons(moving_zscore, nan_index)
    cluster_immobile_neuron = cluster_calcium(data_withNaN, cgIdx_immobile)
    cluster_moving_neuron = cluster_calcium(moving_withNaN, cgIdx_immobile)
    cluster_immobile_zscore = cluster_calcium(immobile_zscore, cgIdx_immobile)
    cluster_moving_zscore = cluster_calcium(moving_zscore, cgIdx_immobile)

    # calculate residuals
    corr_residual = do_residual(clustered_moving_corr, clustered_immobile_corr)
    r_square_residual = do_residual(clustered_moving_r, clustered_immobile_r)
    ava_corr_residual1 = do_residual(AVA_corr1, AVA_corr_im1)
    ava_corr_residual2 = do_residual(AVA_corr2, AVA_corr_im2)
    ava_r_residual1 = do_residual(AVA_r_square1, AVA_r_square_im1)
    ava_r_residual2 = do_residual(AVA_r_square2, AVA_r_square_im2)

    # do the clustering on moving data
    cgIdx_moving = do_clustering(moving_corr_NaN)
    moving_cluster_moving_corr = cluster_matrix(moving_corr_NaN, cgIdx_moving)
    moving_cluster_immobile_corr = cluster_matrix(immobile_corr_NaN, cgIdx_moving)
    moving_cluster_moving_r = cluster_matrix(moving_r_square_NaN, cgIdx_moving)
    moving_cluster_immobile_r = cluster_matrix(immobile_r_NaN, cgIdx_moving)
    AVA_index_mov1 = get_index_AVA_clustered(AVA_neuron1, cgIdx_moving)
    AVA_index_mov2 = get_index_AVA_clustered(AVA_neuron2, cgIdx_moving)

    moving_corr_residual = do_residual(moving_cluster_moving_corr, moving_cluster_immobile_corr)
    moving_r_residual = do_residual(moving_cluster_moving_r, moving_cluster_immobile_r)

    moving_cluster_immobile_neuron = cluster_calcium(data_withNaN, cgIdx_moving)
    moving_cluster_moving_neuron = cluster_calcium(moving_withNaN, cgIdx_moving)
    moving_cluster_immobile_zscore = cluster_calcium(immobile_zscore, cgIdx_moving)
    moving_cluster_moving_zscore = cluster_calcium(moving_zscore, cgIdx_moving)

    AVA_label = [AVA_neuron1, AVA_neuron2] #For dataset 144610
   #PLOTTING

    plot1 = plt.figure(1)
    plt.subplot(1,3,1)
    plt.imshow(clustered_moving_corr, vmin=-1, vmax = 1)
    plt.colorbar()
    plt.title('Correlation_144610_moving')
    plt.xlabel('Sorted Neuron')
    plt.ylabel('Sorted Neuron')
    plt.xticks([AVA_index_imm1,AVA_index_imm2], AVA_label)
    plt.yticks([AVA_index_imm1, AVA_index_imm2], AVA_label)

    plt.subplot(1,3,2)
    plt.imshow(clustered_immobile_corr,vmin=-1, vmax = 1)
    plt.colorbar()
    plt.title('Correlation_144610_immobile')
    plt.xlabel('Sorted Neuron')
    plt.ylabel('Sorted Neuron')
    plt.xticks([AVA_index_imm1,AVA_index_imm2], AVA_label)
    plt.yticks([AVA_index_imm1, AVA_index_imm2], AVA_label)

    plt.subplot(1,3,3)
    plt.imshow(corr_residual, vmin=-1, vmax = 1)
    plt.colorbar()
    plt.title('Residual_immobile-moving_1144610')
    plt.xlabel('Sorted Neuron')
    plt.ylabel('Sorted Neuron')
    plt.xticks([AVA_index_imm1,AVA_index_imm2], AVA_label)
    plt.yticks([AVA_index_imm1, AVA_index_imm2], AVA_label)

    plot2 = plt.figure(2)
    plt.subplot(1,3,1)
    plt.imshow(clustered_moving_r, vmin = 0, vmax = 1)
    plt.colorbar()
    plt.title('R^2_144610_moving')
    plt.xlabel('Sorted Neuron')
    plt.ylabel('Sorted Neuron')
    plt.xticks([AVA_index_imm1,AVA_index_imm2], AVA_label)
    plt.yticks([AVA_index_imm1, AVA_index_imm2], AVA_label)

    plt.subplot(1,3,2)
    plt.imshow(clustered_immobile_r, vmin = 0, vmax = 1)
    plt.colorbar()
    plt.title('R^2_144610_immobile')
    plt.xlabel('Sorted Neuron')
    plt.ylabel('Sorted Neuron')
    plt.xticks([AVA_index_imm1,AVA_index_imm2], AVA_label)
    plt.yticks([AVA_index_imm1, AVA_index_imm2], AVA_label)

    plt.subplot(1,3,3)
    plt.imshow(r_square_residual, vmin=-1, vmax = 1)
    plt.colorbar()
    plt.title('Residual_immobile-moving_1144610')
    plt.xlabel('Sorted Neuron')
    plt.ylabel('Sorted Neuron')
    plt.xticks([AVA_index_imm1,AVA_index_imm2], AVA_label)
    plt.yticks([AVA_index_imm1, AVA_index_imm2], AVA_label)

    plot3 = plt.figure(3)
    plt.subplot(2,3,1)
    plt.bar(neuron_number, AVA_corr1)
    plt.title('AVA Correlation AVA_neuron1_144610_moving')
    plt.xlabel('Neuron')
    plt.ylabel('Correlation')
    plt.ylim(-0.8, 1)
    plt.grid()

    plt.subplot(2,3,2)
    plt.bar(neuron_number, AVA_corr_im1)
    plt.title('AVA Correlation 17_144610_immobile')
    plt.xlabel('Neuron')
    plt.ylabel('Correlation')
    plt.ylim(-0.8, 1)
    plt.grid()

    plt.subplot(2,3,3)
    plt.bar(neuron_number, ava_corr_residual1)
    plt.title('Residual_immobile-moving_1144610')
    plt.xlabel('Neuron')
    plt.ylabel('Residual')
    #plt.ylim(-0.6, 1)
    plt.grid()

    plt.subplot(2,3,4)
    plt.bar(neuron_number, AVA_corr2)
    plt.title('AVA Correlation 81_144610_moving')
    plt.xlabel('Neuron')
    plt.ylabel('Correlation')
    plt.ylim(-0.8, 1)
    plt.grid()

    plt.subplot(2,3,5)
    plt.bar(neuron_number, AVA_corr_im2)
    plt.title('AVA Correlation 81_144610_immobile')
    plt.xlabel('Neuron')
    plt.ylabel('Correlation')
    plt.ylim(-0.8, 1)
    plt.grid()

    plt.subplot(2,3,6)
    plt.bar(neuron_number, ava_corr_residual2)
    plt.title('Residual_immobile-moving_1144610')
    plt.xlabel('Neuron')
    plt.ylabel('Residual')
    #plt.ylim(-0.6, 1)
    plt.grid()

    plot4 = plt.figure(4)
    plt.subplot(2,3,1)
    plt.bar(neuron_number, AVA_r_square1)
    plt.title('AVA R^2 17_144610_moving')
    plt.xlabel('Neuron')
    plt.ylabel('R^2')
    plt.ylim(0, 1)
    plt.grid()

    plt.subplot(2,3,2)
    plt.bar(neuron_number, AVA_r_square_im1)
    plt.title('AVA R^2 17_144610_immobile')
    plt.xlabel('Neuron')
    plt.ylabel('R^2')
    plt.ylim(0, 1)
    plt.grid()

    plt.subplot(2,3,3)
    plt.bar(neuron_number, ava_r_residual1)
    plt.title('Residual_immobile-moving_1144610')
    plt.xlabel('Neuron')
    plt.ylabel('Residual')
    #plt.ylim(0, 1)
    plt.grid()

    plt.subplot(2, 3, 4)
    plt.bar(neuron_number, AVA_r_square2)
    plt.title('AVA R^2 81_144610_moving')
    plt.xlabel('Neuron')
    plt.ylabel('R^2')
    plt.ylim(0, 1)
    plt.grid()

    plt.subplot(2, 3, 5)
    plt.bar(neuron_number, AVA_r_square_im2)
    plt.title('AVA R^2 81_144610_immobile')
    plt.xlabel('Neuron')
    plt.ylabel('R^2')
    plt.ylim(0, 1)
    plt.grid()

    plt.subplot(2, 3, 6)
    plt.bar(neuron_number, ava_r_residual2)
    plt.title('Residual_immobile-moving_1144610')
    plt.xlabel('Neuron')
    plt.ylabel('Residual')
    # plt.ylim(0, 1)
    plt.grid()
    #######

    plot5 = plt.figure(5)
    plt.subplot(1,3,1)
    plt.imshow(moving_cluster_moving_corr, vmin=-1, vmax = 1)
    plt.colorbar()
    plt.title('Correlation_144610_moving')
    plt.xlabel('Sorted Neuron')
    plt.ylabel('Sorted Neuron')
    plt.xticks([AVA_index_mov1,AVA_index_mov2], AVA_label)
    plt.yticks([AVA_index_mov1, AVA_index_mov2], AVA_label)

    plt.subplot(1,3,2)
    plt.imshow(moving_cluster_immobile_corr,vmin=-1, vmax = 1)
    plt.colorbar()
    plt.title('Correlation_144610_immobile')
    plt.xlabel('Sorted Neuron')
    plt.ylabel('Sorted Neuron')
    plt.xticks([AVA_index_mov1,AVA_index_mov2], AVA_label)
    plt.yticks([AVA_index_mov1, AVA_index_mov2], AVA_label)

    plt.subplot(1,3,3)
    plt.imshow(moving_corr_residual, vmin=-1, vmax = 1)
    plt.colorbar()
    plt.title('Residual_immobile-moving_1144610')
    plt.xlabel('Sorted Neuron')
    plt.ylabel('Sorted Neuron')
    plt.xticks([AVA_index_mov1,AVA_index_mov2], AVA_label)
    plt.yticks([AVA_index_mov1, AVA_index_mov2], AVA_label)

    plot6 = plt.figure(6)
    plt.subplot(1,3,1)
    plt.imshow(moving_cluster_moving_r, vmin = 0, vmax = 1)
    plt.colorbar()
    plt.title('R_Square_144610_moving')
    plt.xlabel('Sorted Neuron')
    plt.ylabel('Sorted Neuron')
    plt.xticks([AVA_index_mov1,AVA_index_mov2], AVA_label)
    plt.yticks([AVA_index_mov1, AVA_index_mov2], AVA_label)

    plt.subplot(1,3,2)
    plt.imshow(moving_cluster_immobile_r,vmin = 0, vmax = 1)
    plt.colorbar()
    plt.title('R_Square_144610_immobile')
    plt.xlabel('Sorted Neuron')
    plt.ylabel('Sorted Neuron')
    plt.xticks([AVA_index_mov1,AVA_index_mov2], AVA_label)
    plt.yticks([AVA_index_mov1, AVA_index_mov2], AVA_label)

    plt.subplot(1,3,3)
    plt.imshow(moving_r_residual, vmin=-1, vmax = 1)
    plt.colorbar()
    plt.title('Residual_immobile-moving_1144610')
    plt.xlabel('Sorted Neuron')
    plt.ylabel('Sorted Neuron')
    plt.xticks([AVA_index_mov1,AVA_index_mov2], AVA_label)
    plt.yticks([AVA_index_mov1, AVA_index_mov2], AVA_label)

    plot7 = plt.figure(7)
    plt.subplot(2,2,1)
    plt.imshow(cluster_moving_neuron)
    #plt.colorbar()
    plt.title('Moving_144610_Immobile Cluster')
    plt.xlabel('Time (sec/6)')
    plt.ylabel('Sorted Neuron')

    plt.subplot(2,2,2)
    plt.imshow(cluster_immobile_neuron)
    #plt.colorbar()
    plt.title('Immobile_144610_Immobile Cluster')
    plt.xlabel('Time (sec/6)')
    plt.ylabel('Sorted Neuron')

    plt.subplot(2,2,3)
    plt.imshow(cluster_moving_zscore)
    #plt.colorbar()
    plt.title('Mov_ZScore_144610_Immobile Cluster')
    plt.xlabel('Time (sec/6)')
    plt.ylabel('Sorted Neuron')

    plt.subplot(2,2,4)
    plt.imshow(cluster_immobile_zscore)
    #plt.colorbar()
    plt.title('Imm_Zscore_144610_Immobile Cluster')
    plt.xlabel('Time (sec/6)')
    plt.ylabel('Sorted Neuron')

    plot8 = plt.figure(8)
    plt.subplot(2,2, 1)
    plt.imshow(moving_cluster_moving_neuron)
    #plt.colorbar()
    plt.title('Moving_144610_Moving Cluster')
    plt.xlabel('Time (sec/6)')
    plt.ylabel('Sorted Neuron')

    plt.subplot(2,2, 2)
    plt.imshow(moving_cluster_immobile_neuron)
    #plt.colorbar()
    plt.title('Immobile_144610_Moving Cluster')
    plt.xlabel('Time (sec/6)')
    plt.ylabel('Sorted Neuron')

    plt.subplot(2,2, 3)
    plt.imshow(moving_cluster_moving_zscore)
    #plt.colorbar()
    plt.title('Mov_ZScore_144610_Moving Cluster')
    plt.xlabel('Time (sec/6)')
    plt.ylabel('Sorted Neuron')

    plt.subplot(2, 2,4)
    plt.imshow(moving_cluster_immobile_zscore)
    #plt.colorbar()
    plt.title('Imm_ZScore_144610_Moving Cluster')
    plt.xlabel('Time (sec/6)')
    plt.ylabel('Sorted Neuron')

    plot9 = plt.figure(9)
    plt.subplot(1,3,1)
    plt.imshow(mcluster_move_noNaN, vmin=-1, vmax = 1)
    plt.colorbar()
    plt.title('Move_144610_AvgCorr ='+str(move_noNaN_avg_corr))
    plt.xlabel('Sorted Neuron_NaN removed')
    plt.ylabel('Sorted Neuron_NaN removed')
    plt.xticks(mcluster_AVA, AVA_label)
    plt.yticks(mcluster_AVA, AVA_label)

    plt.subplot(1,3,2)
    plt.imshow(mcluster_imm_noNaN,vmin=-1, vmax = 1)
    plt.colorbar()
    plt.title('Imm_144610_AvgCorr =' +str(imm_nonNaN_avg_corr))
    plt.xlabel('Sorted Neuron_NaN removed')
    plt.ylabel('Sorted Neuron_NaN removed')
    plt.xticks(mcluster_AVA, AVA_label)
    plt.yticks(mcluster_AVA, AVA_label)

    plt.subplot(1,3,3)
    plt.imshow(mcluster_residual, vmin=-1, vmax = 1)
    plt.colorbar()
    plt.title('Residual_immobile-moving_1144610')
    plt.xlabel('Sorted Neuron_NaN removed')
    plt.ylabel('Sorted Neuron_NaN removed')
    plt.xticks(mcluster_AVA, AVA_label)
    plt.yticks(mcluster_AVA, AVA_label)

    plot10 = plt.figure(10)
    plt.subplot(1,3,1)
    plt.imshow(imcluster_move_noNaN, vmin=-1, vmax = 1)
    plt.colorbar()
    plt.title('Move_144610_AvgCorr ='+str(move_noNaN_avg_corr))
    plt.xlabel('Sorted Neuron_NaN removed')
    plt.ylabel('Sorted Neuron_NaN removed')
    plt.xticks(imcluster_AVA, AVA_label)
    plt.yticks(imcluster_AVA, AVA_label)

    plt.subplot(1,3,2)
    plt.imshow(imcluster_imm_noNaN,vmin=-1, vmax = 1)
    plt.colorbar()
    plt.title('Imm_144610_AvgCorr =' +str(imm_nonNaN_avg_corr))
    plt.xlabel('Sorted Neuron_NaN removed')
    plt.ylabel('Sorted Neuron_NaN removed')
    plt.xticks(imcluster_AVA, AVA_label)
    plt.yticks(imcluster_AVA, AVA_label)

    plt.subplot(1,3,3)
    plt.imshow(imcluster_residual, vmin=-1, vmax = 1)
    plt.colorbar()
    plt.title('Residual_immobile-moving_1144610')
    plt.xlabel('Sorted Neuron_NaN removed')
    plt.ylabel('Sorted Neuron_NaN removed')
    plt.xticks(imcluster_AVA, AVA_label)
    plt.yticks(imcluster_AVA, AVA_label)

    plot11 = plt.figure(11)
    plt.subplot(3,1,1)
    plt.plot(moving_velocity) #, moving_time)
    plt.title('Velocity')
    plt.xlabel('time')

    plt.subplot(3,1,2)
    plt.plot(moving_data[AVA_neuron1,:]) #, moving_time)
    plt.title('AVA Neuron 17')
    plt.xlabel('time_contig')

    plt.subplot(3,1,3)
    plt.plot(moving_data[AVA_neuron2,:])#, moving_time)
    plt.title('AVA Neuron 81')
    plt.xlabel('time_contig')
    plt.show()


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Oct 20 12:26:44 2019

@author: fubao
"""

# facebook minesweeper

'''
529. Minesweeper
Medium

387

373

Favorite

Share
Let's play the minesweeper game (Wikipedia, online game)!

You are given a 2D char matrix representing the game board. 'M' represents an unrevealed mine, 'E' represents an unrevealed empty square, 'B' represents a revealed blank square that has no adjacent (above, below, left, right, and all 4 diagonals) mines, digit ('1' to '8') represents how many mines are adjacent to this revealed square, and finally 'X' represents a revealed mine.

Now given the next click position (row and column indices) among all the unrevealed squares ('M' or 'E'), return the board after revealing this position according to the following rules:

If a mine ('M') is revealed, then the game is over - change it to 'X'.
If an empty square ('E') with no adjacent mines is revealed, then change it to revealed blank ('B') and all of its adjacent unrevealed squares should be revealed recursively.
If an empty square ('E') with at least one adjacent mine is revealed, then change it to a digit ('1' to '8') representing the number of adjacent mines.
Return the board when no more squares will be revealed.

Example 1:

Input:

[['E', 'E', 'E', 'E', 'E'],
 ['E', 'E', 'M', 'E', 'E'],
 ['E', 'E', 'E', 'E', 'E'],
 ['E', 'E', 'E', 'E', 'E']]

Click : [3,0]

Output:

[['B', '1', 'E', '1', 'B'],
 ['B', '1', 'M', '1', 'B'],
 ['B', '1', '1', '1', 'B'],
 ['B', 'B', 'B', 'B', 'B']]

Explanation:

Example 2:

Input:

[['B', '1', 'E', '1', 'B'],
 ['B', '1', 'M', '1', 'B'],
 ['B', '1', '1', '1', 'B'],
 ['B', 'B', 'B', 'B', 'B']]

Click : [1,2]

Output:

[['B', '1', 'E', '1', 'B'],
 ['B', '1', 'X', '1', 'B'],
 ['B', '1', '1', '1', 'B'],
 ['B', 'B', 'B', 'B', 'B']]

Explanation:

'''

class Solution:
    def updateBoard(self, board: List[List[str]], click: List[int]) -> List[List[str]]:
        #SUCCESS case 1: if current pos is M, then assign it with X
        # case 2: if all 8 adjacent block has M, then count the number of M for each position if the count is not 0, assign the count to current pos; if te count is 0,then assin the pos with "B", then wee continue to find all its 8 adjacent blocks recursively.
        if board is None or len(board) == 0 or len(board[0]) == 0:
            return board

        m = len(board)
        n = len(board[0])

        row = click[0]
        col = click[1]

        cnt = 0    # how many mines in the adjacent neighbors
        if board[row][col] == 'M':
            board[row][col] = 'X'

        else:
            # check all it 8 adjacent
            for i in range(-1, 2):
                for j in range(-1,2):
                    x = row + i
                    y = col + j
                    # out of boundary
                    if x < 0 or x >= m or y < 0 or y >= n:
                        continue
                    if board[x][y] == 'M':
                        cnt += 1
            if cnt > 0:
                board[row][col] = str(cnt)
            else:
                board[row][col] = 'B'
                for i in range(-1, 2):
                    for j in range(-1, 2):
                        x = row + i
                        y = col + j
                        # out of boundary
                        if x < 0 or x >= m or y < 0 or y >= n:
                            continue
                        if board[x][y] == 'E':
                            self.updateBoard(board, [x,y])

        return board

        # optimize the code above
        # https://www.cnblogs.com/grandyang/p/6536694.html


# -*- coding: utf-8 -*-
"""
Created on Sun Jun 26 20:46:50 2016

@author: phyrct
"""
import numpy as np
from kermit import KerrList, KerrArray

def hysteresis(images, fieldlist=None, box=None):
    """Make a hysteresis loop of the average intensity in the given images

    Parameters
    ----------
    ims: Kerrlist
        list of images for extracting hysteresis
    fieldlist: list or tuple
        list of fields used, if None it will try to get field from imgae metadata
    box: list
        [xmin,xmax,ymin,ymax] region of interest for hysteresis

    Returns
    -------
    hyst: nx2 np.ndarray
        fields, intensities, 2 column numpy array
    """
    if not isinstance(images,KerrList):
        raise TypeError('images must be a KerrList')
    hys_length=len(images)
    if fieldlist is None:
        fieldlist=images.slice_metadata(key='Field', values_only=True)
    fieldlist=np.array(fieldlist)
    assert len(fieldlist)==hys_length, 'images and field list must be of equal length'
    assert all([i.shape==images[0].shape for i in images]), 'images must all be same shape'
    if box is None:
        box=(0,images[0].shape[1],0,images[0].shape[0])

    hyst=np.column_stack((fieldlist,np.zeros(hys_length)))
    for i,im in enumerate(images):
        im=im.crop_image(box=box, copy=True)
        hyst[i,1] = np.average(im)
    return hyst

def correct_drift(images, refindex, threshold=0.005, upsample_factor=50):
    """Align images to correct for image drift."""
    ret=images.apply_all('correct_drift', ref=images[refindex],
                                                threshold=threshold)
    return images


# -*- coding: utf-8 -*-
"""
Additional requests
"""
from simple_settings import settings
from common_modules.parser import ParserPost, Parser

__author__ = 'Vadim Arsenev'
__version__ = '1.1.0'
__data__ = '04.08.2021'

def add_html(treeid):
    params = {'treeId': treeid, 'siteStyle': 'SIMPLE'}
    tournamentMarket = Parser(settings.MARKETS_URL, params=params)
    html = tournamentMarket.parser_resultHtml()

    return html

def add_html_old(treeid):
    payload = {'treeId': treeid, 'columnSize': '8', 'siteStyle': 'SIMPLE'}
    tournamentMarket = ParserPost(settings.MARKETS_URL, payload)
    html = tournamentMarket.parser_resultHtml()

    return html


n= int(input())

A=list(map(int,input().split()))
AA=sorted(A)
p=[]
for i in A:
    p.append(AA.index(i))
    AA[AA.index(i)]= -1

for j in p:
    print(j,end= " ")


from django.urls import path

from . import views

urlpatterns = [
    path('', views.PostListView.as_view(), name='posts'),
    path('search/', views.search_post, name='search'),
    path('post/<int:post_id>/', views.post_detail, name='post_detail'),
    path('post/<int:post_id>/add_comment/', views.add_comment, name='add_comment'),
    path('add_post/', views.add_post, name='add_post'),
]


"""
This software is intended to bulk download maps from
https://www.opengeodata.nrw.de/produkte/geobasis/tk/

"""

import argparse
import json
import requests
import tempfile
import os

"""
The version-number
"""
version = (0, 0, 2)
__version__ = '.'.join('%d' % p for p in version)

"""
A dict of product names and URLs
"""
__products__ = {
    'dtk25': 'https://www.opengeodata.nrw.de/produkte/geobasis/tk/tk25/tk25_s_pdf/',
    'dtk50': 'https://www.opengeodata.nrw.de/produkte/geobasis/tk/tk50/tk50_s_pdf/',
    'dtk100': 'https://www.opengeodata.nrw.de/produkte/geobasis/tk/tk100/tk100_s_pdf/',
}

class DownloadError(Exception):
    pass

def http_downloader(url: str, mode: str=''):
    """

    Args:
        url: The document which shall be downloaded
        mode: The mode in which the download should take place, can be 'json', defaults to ''

    Returns: The content of the download either as parsed json when mode was 'json' or otherwise.

    """
    r = requests.get(url,  timeout=10)

    if r.status_code != requests.codes.OK:
        # something went wrong.
        raise DownloadError('Download error. Getting %s returned %s' % (url, r.status_code))

    if mode == 'json':
        return r.json
    if mode == 'text':
        return r.text
    else:
        return r.content

def get_files_from_index_json(index: str):
    """
    Reads the index-file of the map-collection and parses it into a dict.

    Args:
        index: the json-document

    Returns: An array containing dicts which describe the downloadable file.
             The description consist of: name, size, timestamp

    """
    index_dl = http_downloader(index, 'text')
    j_index = json.loads(index_dl)
    datasets = j_index.get('datasets', None)
    metafiles = j_index.get('metafiles', None)

    if datasets == None:
        raise ValueError('The index-document does not contain datasets.')

    files = []

    for ds in datasets:
        # iterate the datasets and add all files to the files
        # this iteration is done, in hope that it works if multiple datasets
        # occur. With one dataset the result should be the same as datasets[0].get('files')
        fs = ds.get('files', None)
        if fs:
            for f in fs:
                files.append(f)

    if metafiles:
        for ms in metafiles:
            # Iterate the list of metafiles in the index.json and append files.
            files.append(ms)

    if files is None:
        raise ValueError('The index-document does not contain files.')

    return files

def get_product_url(product: str):
    """
    Args:
        product: The product which should be downloaded, dtk50 or dtk100

    Returns: the URL which is stored in the dict above.

    """
    url = __products__.get(product, None)
    if url is None:
        raise NotImplementedError('Product error: %s is unkown' % (product))
    return url

def create_download_folder(product: str, path: str=None):
    """

    Args:
        product: The name of the product
        path: The path where the download-folder shall be created, defaults to temp

    Returns: The path to the download-folder

    """

    if path is None:
        abs_path_to_folder = tempfile.mkdtemp(suffix=None, prefix=product + '_')
    else:
        raise NotImplementedError('Creating download Directories outside of temp is not supported.')

    return abs_path_to_folder

def file_downloader(files, producturl: str, downloaddir: str):
    """
    Args:
        files: An array containing the file-names
        producturl: The URL of the products
        downloaddir: The direcory where the new files shall be created.

    Returns: None

    """
    # TODO: Replace Print-Outs with a Logger

    print("Downloading %s files." % (files.__len__()))
    dl_ctr = 0
    for f in files:
        fname = f.get('name', None)
        if not fname:
            # Skip files without a name
            continue

        fileurl = producturl + fname
        # now download
        print("Downloading %s." % (fileurl))

        content = http_downloader(fileurl)
        with open(os.path.join(downloaddir, fname), 'w+b') as outfile:
            outfile.write(content)

        dl_ctr+=1
        print("Downloaded %s of %s files." % (dl_ctr , files.__len__()))

    return None

def main():
    # TODO: Replace Print-Outs with a Logger

    parser = argparse.ArgumentParser(description='Bulk-Download DTK-Maps as PDFs from https://www.opengeodata.nrw.de')

    parser.add_argument('--version', action='version', version=('Version %s' % __version__))
    parser.add_argument('--product', nargs='?', choices=__products__.keys(), default='dtk50')

    args = parser.parse_args()

    product = args.product
    product_url = get_product_url(product)

    # Step 1: Parse the Index file into "files"
    files = get_files_from_index_json(product_url + 'index.json')

    # Step 2: Create the Download-Folder
    dlpath = create_download_folder(product)
    print('Data will be downloaded to: %s' % dlpath)

    # Step 3: Iterate "files" and download every file an write it to the Download-folder
    file_downloader(files, product_url, dlpath)

    exit(0)

if __name__ == "__main__":
    main()


from django.urls import path
from inv import views

#######################alex 2020/09/23 app 登录##############################
urlpatterns = [
    path('', views.Productlist.as_view()),                 #unit list
    path('<int:pk>', views.ProductDetail.as_view()),     #unit Detail
    path('stock',views.Stocklist.as_view()),                 #category list
    path('stock/<int:pk>', views.StockDetail.as_view()),      #category Detail
    path('stock/out', views.Out_of_the_warehouse.as_view()),      #out of the warehouse
    path('stock/enter', views.Enter_of_the_warehouse.as_view()),      #enter of the warehouse
    path('stock/product/<int:pk>', views.Current_quantity.as_view()), # current_quantity product
]

############################################################################


import xml.etree.ElementTree as ET
import os

directory = input("Enter directory: ")

# Get a list of all .xml files in the directory
xml_files = [f for f in os.listdir(directory) if f.endswith(".xml")]

# The properties that you want to add to your .xml files
properties = """
<properties>
  <prop name="Name">
   <prop name="es_MX.utf8" />
  </prop>
  <prop name="Size">2085 796</prop>
  <prop name="BackColor">ISA_11_BackgroundGrey</prop>
 </properties>
"""

# Counter for the number of modified files
modified_files = 0

for xml_file in xml_files:
  # Check if the file ends with "SYM.xml"
  if xml_file.endswith("DETAILED.xml"):
    tree = ET.parse(os.path.join(directory, xml_file))
    root = tree.getroot()

    properties_elem = root.find("properties")

    if properties_elem is not None:
      root.remove(properties_elem)
      root.append(ET.fromstring(properties))

    tree.write(os.path.join(directory, xml_file))

    modified_files += 1

    print(f"Modified file: {xml_file}")

print(f"Total modified files: {modified_files}")


import cv2
import os

def video_to_image(image_folder, video_name):
    cap = cv2.VideoCapture(video_name)
    fps = int(cap.get(cv2.CAP_PROP_FPS))

    frame_count = 0

    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frame_count += 1
        image_name = f"frame_{frame_count:04d}.jpg"
        image_path = os.path.join(image_folder, image_name)

        cv2.imwrite(image_path, frame)

    cap.release()
    cv2.destroyAllWindows()

# 使用方法示例
image_folder = 'D:/ultralytics-main/pictures'  # 图片文件夹路径
video_name = 'D:/微信文件/WeChat Files/wxid_jmmnslqz5lwl22/FileStorage/File/2023-08/video.avi'  # 输出视频文件名

video_to_image(image_folder, video_name)


__docformat__ = "google"
# build doc : pdoc -o D:\Dropbox\Run4_transitoire\Documentation D:\Dropbox\Run4_transitoire\Functions
# add project folder path to Anaconda\Lib\site-packages in python38.pth

"""
    This module holds the code to prepare data and interact with Feflow through the official ifm38 library
"""

import gc
import time
from os import chdir, listdir, remove
import os.path
from os.path import isfile, join

import pandas as pd
import numpy as np

# Path to the FeFlow binaries has to bet set before importing ifm38
#chdir(r'C:\\Program Files\\DHI\\2021\\FEFLOW 7.4\\bin64')

#   To import ifm38 from anaconda prompt, it's needed to add in the folder
#   Anaconda\Lib\site-packages a file named python38.pth with the direct path
#   to feflow, in my case : C:\\Program Files\\DHI\\2021\\FEFLOW 7.4\\bin64
#import ifm38 as ifm

""" ------------------------- Defining Feflow Callbacks Function to act on simulator ------------------------- """

def preEnterSimulator(doc):

    """
    Assigning parameters to Feflow before launching any simulation. Parameters are passed
    by global variables has defined in Feflow documentation. In this code, we send permKx,
    permKz and recharge issued from geostatistical simulations.

    Args:
        doc (PyFeflowDoc): Feflow Document imported with ifm.loadDocument(fem_file)

    Raises:
        ValueError: Wrong number of elements in _PERM_KX
        ValueError: Wrong number of elements in _PERM_KY
        ValueError: Wrong number of elements in _PERM_KZ
        ValueError: Wrong number of recharge in _RECHARGE
        ValueError: Wrong number of nodes while importing steady heads
    """

    if (_ELEMENT_NUMBER!=len(_PERM_KX)):
        raise ValueError("Wrong number of elements in permKx")

    # To modify when switching to full anisotropic k fields
    if (_ELEMENT_NUMBER!=len(_PERM_KX)):
        raise ValueError("Wrong number of elements in permKy")

    if (_ELEMENT_NUMBER!=len(_PERM_KZ)):
        raise ValueError("Wrong number of elements in permKz")

    #if (len(doc.getParamValues(ifm.Enum.P_IOFLOW))!=len(_RECHARGE)):
    #    raise ValueError("Wrong number of recharge zones")

    # In this case, Kx = Ky, _PERM_KX is assigned to P_CONDX and P_CONDY

    doc.setParamValues(ifm.Enum.P_CONDX,_PERM_KX,0)
    doc.setParamValues(ifm.Enum.P_CONDY,_PERM_KX,0)
    doc.setParamValues(ifm.Enum.P_CONDZ,_PERM_KZ,0)

    #doc.setParamValues(ifm.Enum.P_IOFLOW,_RECHARGE,0)
    doc.setParamValues(ifm.Enum.P_COMP,_SS,0)

    # In transient case, we want to set hydraulic head simulated in steady state simulation
    # as initial hydraulic heads to speed up simulation time.

    if _TRANSIENT == True :

        #steady_heads = np.load(_PATH_SAVE+"/temp/steady_heads_"+str(_MODEL_NUMBER)+".npy")
        steady_heads = _TEMP_HEAD

        if (doc.getNumberOfNodes()!=len(steady_heads)):
            raise ValueError("Wrong number of nodes while importing steady heads")

        #doc.setParamValues(ifm.Enum.P_HEAD,steady_heads.tolist(),0)
        doc.setParamValues(ifm.Enum.P_HEAD,steady_heads,0)

        del steady_heads

def postTimeStep(doc):

    """
    Assigning parameters to Feflow before launching any simulation. Parameters are passed
    by global variables has defined in Feflow documentation. In this code, we send permKx,
    permKz and recharge issued from geostatistical simulations.

    Args:
        doc (PyFeflowDoc): Feflow Document imported with ifm.loadDocument(fem_file)
    """

    # In steady state we stored the hydraulic heads simulated at each nodes to use them
    # as initial hydraulic head in transient simulation
    if _STEADY == True :
        global _TEMP_HEAD
        steady_heads = doc.getParamValues(ifm.Enum.P_HEAD)

        #np.save(_PATH_SAVE+"/temp/steady_heads_"+str(_MODEL_NUMBER),steady_heads)
        _TEMP_HEAD = steady_heads

        array_sum = np.sum(steady_heads)
        array_has_nan = np.isnan(array_sum)

        if np.max(steady_heads) > 100 or np.min(steady_heads) < 1 or array_has_nan :
            global _STEADY_CONVERGENCE
            _STEADY_CONVERGENCE = False

        del steady_heads

        global _STEADY_HEAD
        temp_heads = np.zeros(len(_WELL))

        for well in range(len(_WELL)) :
            temp_heads[well] = doc.getResultsFlowHeadValueAtXYZ(_WELL[well][0], _WELL[well][1], _WELL[well][2])

        _STEADY_HEAD.append(temp_heads)

    else :

        # well + 1 to store an additionnal columns with time of simulation
        temp_heads = np.zeros(len(_WELL)+1)

        _time = doc.getAbsoluteSimulationTime()

        print("Time step :"+str(np.round(_time,1))+"d, Total ellapsed time : "+str(int(time.time()-_TIME))+"s")

        global _HEAD
        for well in range(len(_WELL)) :
            temp_heads[0] = _time
            temp_heads[well+1] = doc.getResultsFlowHeadValueAtXYZ(_WELL[well][0], _WELL[well][1], _WELL[well][2])
        _HEAD.append(temp_heads)

        # Put a treshold to delete non converging simulations exceeding a given time (mean converging time on succesfull simulation + 200 sec buffer)
        if (_time < 3 and (time.time()-_TIME)> 550) or (_time < 6 and (time.time()-_TIME)> 800) or (_time < 11 and (time.time()-_TIME)> 1000) or (_time < 20 and (time.time()-_TIME)> 1200) or (time.time()-_TIME> 1600) :

            doc.stopSimulator()
            doc.closeDocument()
            global _TRANSIENT_CONVERGENCE
            _TRANSIENT_CONVERGENCE = False

def check_launch_feflow(path_data: str,
                    fem_steady_model: str,
                    fem_transient_model: str,
                    assim_number: int,
                    model_number: int) -> None:

    """
    To avoid redoing lengthy simulation on already done one, we check the list of simulation to do
    versus the one already done. We also checked the existence of folders.

    Args:
        path_data (str): Abolute path to the Data folder contained in the project
        fem_steady_model (str): Abolute path to the Feflow steady state model
        fem_transient_model (str): Abolute path to the Feflow transient state model
        assim_number (int): Actual assimilation step in the multiple data assimilation process
        model_number (int): Number of the model to be run in the actual simulation

    Raises:
        ValueError: Individual_Update folder missing in Data/Assim_X/ folder
        ValueError: Head folder missing in Data/Assim_X/ folder
    """

    str_model_number = str(model_number)
    str_assim_number = str(assim_number)

    temp_filename = "assim_"+str_assim_number+"_model_"+str_model_number+".npy"

    path = path_data + "Assim_"+str_assim_number+"/Individual_Update/"
    path_already_done = path_data + "Assim_"+str_assim_number+"/Head/"

    if not os.path.isdir(path):
        raise ValueError("Individual_Update folder missing in Assim_"+str_assim_number+"/Individual_Update/")

    if not os.path.isdir(path_already_done):
        raise ValueError("Head folder missing in Assim_"+str_assim_number+"/Head/")

    files = [f for f in listdir(path) if isfile(join(path, f))]
    files_already_done = [f for f in listdir(path_already_done) if isfile(join(path_already_done, f))]

    if "head_"+str(model_number)+".npy" in files_already_done :
        print("Loading succesfull, permeability field "+str_model_number+" already run")

    else :

        temp_filename = "assim_"+str_assim_number+"_model_"+str_model_number+".npy"

        if temp_filename in files :

            print("Loading succesfull, permeability field "+str_model_number+" being run")

            launch_feflow(path_data,
                                fem_steady_model,
                                fem_transient_model,
                                assim_number,
                                model_number)

        else :
            print("/!\ /!\ Loading unsuccesfull, permeability field assim_"+str_assim_number+"_model_"+\
                str_model_number+".npy is missing.\n Could be normal in case \
            of dropped non converging model /!\ /!\\")

def launch_feflow(path_data: str,
                    fem_steady_model: str,
                    fem_transient_model: str,
                    assim_number: int,
                    model: int) -> None:
    """
    Functions to assign parameters to Feflow and run the simulation directly from Python.
    A lot of internal global variable are defined in accordance with Feflow documentation.
    They follow a naming scheme to minimize risk of interference with other python file :
    _NAME_OF_INTERNAL_GLOBAL.

    First step is to assign parameters to the steady state Feflow model and running it. Hydraulic
    heads results are then stored as initial condition from the transient model run afterward.

    Args:
        path_data (str): Abolute path to the Data folder contained in the project
        fem_steady_model (str): Abolute path to the Feflow steady state model
        fem_transient_model (str): Abolute path to the Feflow transient state model
        assim_number (int): Actual assimilation step in the multiple data assimilation process
        model_number (int): Number of the model to be run in the actual simulation

    Raises:
        ValueError: Mising steady state Feflow model, check the name or file
        ValueError: Mising transient state Feflow model, check the name or file
    """

    if not os.path.isfile(fem_steady_model):
        raise ValueError("Mising steady state Feflow model, check the name or file")

    if not os.path.isfile(fem_transient_model):
        raise ValueError("Mising transient state Feflow model, check the name or file")

    global _ELEMENT_NUMBER
    global _K_ZONE
    global _RECHARGE_ZONE
    global _UPDATED_ELEMENT
    global _PERM_KX
    global _PERM_KZ
    global _RECHARGE
    global _SS
    global _PATH
    global _PATH_SAVE
    global _FEM_FILE
    global _STEADY
    global _TRANSIENT
    global _HEAD
    global _STEADY_HEAD
    global _WELL
    global _TIME
    global _MODEL_NUMBER
    global _STEADY_CONVERGENCE
    global _TRANSIENT_CONVERGENCE
    global _TEMP_HEAD

    _MODEL_NUMBER = model

    _TIME = time.time()

    # Loading all relevant files, Feflow model, updated parameters, and the kZone and recharge zone
    # linking the global parameters to the element number in the Feflow model.
    _FEM_FILE = fem_steady_model
    doc = ifm.loadDocument(_FEM_FILE)

    _ELEMENT_NUMBER = doc.getNumberOfElements()

    _UPDATED_ELEMENT = np.load(path_data+"/FEFLOW_Model/updated_element.npy")

    # Elements number count starts at 1 in feflow, check the consitency with input updated_element.npy indexing
    _UPDATED_ELEMENT = (_UPDATED_ELEMENT - 1).astype(int)

    _PATH = path_data+"Assim_"+str(assim_number)+"/Individual_Update/"
    _PATH_SAVE = path_data+"Assim_"+str(assim_number)+"/"

    # _HEAD is fullfilled in the postTimeStep() function called by Feflow
    _HEAD = []
    _STEADY_HEAD = []
    _TEMP_HEAD = []

    # Coordinates of all observations point used in the model for assimilation
    # and calibration purpose. They are monitored at each time of simulation.
    _WELL = pd.read_csv(path_data+"/FEFLOW_Model/well_location.inp")

    _WELL = np.vstack((_WELL.X.values-doc.getOriginX(),
                       _WELL.Y.values-doc.getOriginY(),
                       _WELL.Z.values-doc.getOriginZ())).T

    # updated_parameter are coming from the assimilation process, in this model, all data except
    # the last 47 are local element, with a value for each cell in the model. They are called
    # local_updated_parameter for this reason.
    # The last 47 parameters assimilated are global parameters, influencing several cells in
    # Feflow model through kZones. They are treated separately to bring back kZones to elements.
    # On the 45 parameters, the first 15 are related to kx, the next 15 to kr (ratio of kz to kx),
    # the last 15 are saturation coefficient in different surface zones.

    elementsCoordinates = pd.read_pickle(path_data+"/FEFLOW_Model/feflow_element")

    _K_ZONE = elementsCoordinates.K_zone.values

    parameter_name = np.unique(elementsCoordinates.K_zone.values)

    updated_parameter = np.load(path_data+"Assim_"+str(assim_number)+"/Individual_Update/assim_"+\
            str(assim_number)+"_model_"+str(model)+".npy")

    global_parameters = len(updated_parameter) - _ELEMENT_NUMBER

    local_updated_parameter = updated_parameter[:-global_parameters]
    global_updated_parameters = updated_parameter[-global_parameters:]

    kx = np.array(doc.getParamValues(ifm.Enum.P_CONDX))

    """
    We will let the kx value set outside of the local update

    # First we applied permeability coming from global variable, by applying global
    # value to all elements, local elements will replace these global values where
    # they are defined.
    i = 0
    for name in parameter_name :
        # Feflow takes permeabilty as m/d, in case of m/s, multiplying by 3600 * 24.
        kx[np.where(_K_ZONE==name)] = global_updated_parameters[i] *3600*24
        i +=1

    """

    kx[_UPDATED_ELEMENT] = local_updated_parameter[_UPDATED_ELEMENT]

    ### ky = kx, no need to add another vector in memory

    kz_temp = np.zeros(_ELEMENT_NUMBER)
    kz = np.array(doc.getParamValues(ifm.Enum.P_CONDZ))

    # First we applied permeability coming from global variable, by applying global
    # value to all elements, local elements will replace these global values where
    # they are defined.
    i = 15
    for name in parameter_name :
        # Same logic, but kz is defined as a ratio of conductivity to kx.
        kz_temp[np.where(_K_ZONE==name)] = kx[np.where(_K_ZONE==name)] / global_updated_parameters[i]
        i += 1

    kz[_UPDATED_ELEMENT] = kz_temp[_UPDATED_ELEMENT]

    ss_temp = np.zeros(_ELEMENT_NUMBER)
    ss = np.array(doc.getParamValues(ifm.Enum.P_COMP))
    i = 30
    for name in parameter_name :
        ss_temp[np.where(_K_ZONE==name)] =  global_updated_parameters[i]
        i += 1

    ss[_UPDATED_ELEMENT] = ss_temp[_UPDATED_ELEMENT]

    # Feflow is unable to deal with direct numpy array, he needs list
    _PERM_KX = kx.tolist()
    _PERM_KZ = kz.tolist()
    _SS = ss.tolist()

    # To reduce memory impact in case of big Feflow model, we free memory as much as possible
    del _K_ZONE
    del _UPDATED_ELEMENT
    del kx
    del kz
    del kz_temp
    del ss
    del ss_temp
    #del kr
    del updated_parameter
    del local_updated_parameter
    del global_updated_parameters
    del parameter_name

    # We need to specify in which kind of simulation we are
    _STEADY = True
    _TRANSIENT = False
    _STEADY_CONVERGENCE = True

    doc.startSimulator()
    doc.stopSimulator()

    doc.saveDocument(fem_steady_model[:-4]+"_temp.fem")

    doc.closeDocument()

    print("Steady simulation "+str(model)+" done in "+str(int(time.time()-_TIME))+"s")

    temp_time = _TIME
    # Reset timer for transient simulation
    _TIME = time.time()

    if _STEADY_CONVERGENCE :
        print("Model has converged")
        _FEM_FILE = fem_transient_model

        _STEADY_HEAD = np.array(_STEADY_HEAD)

        # Steady Head used to debug, to check the convergence of the model before next step
        np.save(_PATH_SAVE+"Steady_head/steady_head_"+str(model)+'.npy',_STEADY_HEAD)

        doc = ifm.loadDocument(_FEM_FILE)

        _STEADY = False
        _TRANSIENT = True
        _TRANSIENT_CONVERGENCE = True

        doc.startSimulator()

        if _TRANSIENT_CONVERGENCE :
            doc.stopSimulator()

            doc.closeDocument()

            _HEAD = np.array(_HEAD)

            np.save(_PATH_SAVE+"Head/head_"+str(model)+'.npy',_HEAD)

        else :
            print("Convergence failed with this permeability field")
            print("Deletion of the permeability field")

            os.remove(path_data+"Assim_"+str(assim_number)+"/Individual_Update/assim_"+\
                str(assim_number)+"_model_"+str(model)+".npy")

    else :
        print("Convergence failed with this permeability field")
        print("Deletion of the permeability field")

        os.remove(path_data+"Assim_"+str(assim_number)+"/Individual_Update/assim_"+\
            str(assim_number)+"_model_"+str(model)+".npy")

    del _HEAD
    del _WELL
    del _STEADY
    del _TRANSIENT
    del _ELEMENT_NUMBER
    del _TEMP_HEAD

    # Using a try before deleting these vectors, if Feflow hasn't deallocating them properly
    try:
        del _PERM_KX
    except :
        pass

    try:
        del _PERM_KZ
    except :
        pass

    try:
        del _SS
    except :
        pass
    # try:
    #     del _RECHARGE
    # except :
    #     pass

    print("Total simulation "+str(model)+" done in "+str(int((time.time()-temp_time)//60))+" minutes")

    del _TIME

    gc.collect()