"""
Console output colorizer. pyreadline required( on windows )
"""
__author__ = "Jussi Toivola"
__license__ = "MIT License"

from SCons.Platform import win32, posix
import os
import subprocess as sp
import sys

#: Pyreadline console
CONSOLE = None
try:
    if os.name == "nt":
        import pyreadline
        CONSOLE = pyreadline.GetOutputFile()
except ImportError: #IGNORE:W0704
    pass

#: Color constants. TODO: Get from user settings.
class Colors( object ):
    WARNING = 14
    ERROR = 12
    NORMAL = 7
    COMMENT = 10

if sys.platform == "linux2":
    Colors.WARNING = 33
    Colors.ERROR = 31
    Colors.NORMAL = 0
    Colors.COMMENT = 32

#: Get initial color of the console
if CONSOLE is not None:
    Colors.NORMAL = CONSOLE.attr

#: Lines containing these strings are drawn in Colors.ERROR color
KEYWORD_ERRORS = [ __x.lower() for __x in [ "error:", " error ", "failed", " undefined ", "does not match", "illegal",
               "was expected", "The process cannot access",
               "No such file or directory", "*** ", "cannot be opened",
               "Cannot convert", "needed by", "explicit dependency" ] ]
KEYWORD_ERRORS.sort()

#: Excluded errors
KEYWORD_ERRORS_EXCLUDE = [ "error.c", "error.cpp", "error.h"  ]

#: Drawn with Colors.COMMENT
KEYWORD_COMMENT = [ __x.lower() for __x in [ "scons:", "note:", "copy(", "install file"] ]

#: Excluded comments
KEYWORD_COMMENT_EXCLUDE = []

#: Warning keywords. Drawn with Colors.WARNING
KEYWORD_WARNINGS = ["warning ", "warning:"]

#: Excluded warnings
KEYWORD_WARNINGS_EXCLUDE = ["-warning"]

#: Map colors to keywords.
KEYWORDMAP = [ ( KEYWORD_WARNINGS, KEYWORD_WARNINGS_EXCLUDE, Colors.WARNING ),
               ( KEYWORD_ERRORS, KEYWORD_ERRORS_EXCLUDE, Colors.ERROR ),
               ( KEYWORD_COMMENT, KEYWORD_COMMENT_EXCLUDE, Colors.COMMENT ) ]

class LineCounts( object ):
    def __init__( self ):
        self.Errors = 0
        self.Warnings = 0
        self.Others = 0
        self.reset()

    def reset( self ):
        self.Errors = 0
        self.Warnings = 0
        self.Others = 0

LINECOUNTS = LineCounts()

#: Original sys.stdout
savedstdout = sys.stdout
savedstderr = sys.stderr

def subsitute_env_vars( line, env ):
    """ Substitutes environment variables in the command line.

    E.g. dir %EPOCROOT% -> dir T:\

    The command line does not seem to do this automatically when launched
    via subprocess.
    """
    for key, value in env.items():
        line = line.replace( "%%%s%%" % key, value )
    return line

def _handle_posix_write(line,color):
    msg = "%02i" % color
    savedstdout.write( "\x1b[%sm" % ( msg ) )
    savedstdout.write( line )

    # Reset
    msg = "%02i" % 0
    savedstdout.write( "\x1b[%sm" % ( msg ) )

def write( line, color ):
    """Write line with color"""
    global CONSOLE

    if os.name == "posix" and sys.__stdout__.isatty():
        return _handle_posix_write( line, color )

    elif CONSOLE is not None:
        try:
            CONSOLE.write_color( line, color )
            CONSOLE.write_color( "", CONSOLE.attr )
            return

        except TypeError:
            # Occurs at least when redirecting data to a file
            # So disable the color support
            CONSOLE = None

    if color in [ Colors.ERROR, Colors.WARNING ]:
        savedstderr.write( line )
    else:
        savedstdout.write( line )

def comment( line ):
    "Util for comments"
    write( line + "\n", Colors.COMMENT )

def error( line ):
    "Util for errors"
    write( line, Colors.ERROR )

class ConsoleBase( object ):
    """Base class for colored console output"""

    def __init__( self ):
        self.savedstdout = sys.stdout
        self.savedstderr = sys.stderr

    def flush( self ):
        self.savedstderr.flush()
        self.savedstdout.flush()

    def read( self ):
        """Read both stdout and stderr and output the data with colors"""
        block = None
        result = ""

        while block != "":
            block = self.savedstdout.read()
            self.write( block )
            result += block

            block = self.savedstderr.read()
            self.write( block )
            result += block

        return result

    def write( self, text ):
        """Colorize each line based on the keywords"""

        lines = text.split( "\n" )

        for i in xrange( len( lines ) ):

            line = lines[i]
            lowcase = line.lower()

            detected = False
            for kws, excludes, color in KEYWORDMAP:
                # Check keywords and exclusion keywords
                if  len( [x for x in kws if x in lowcase] ) > 0 \
                and len( [x for x in excludes if x in lowcase] ) == 0:
                    write( line, color )
                    detected = True
                    break

            # Draw in normal color then
            if not detected:
                write( line, Colors.NORMAL )

            if i < ( len( lines ) - 1 ):
                write( "\n", Colors.NORMAL )
class OutputConsole( ConsoleBase ):
    """Handles spawning external processes and colors the console output.
    """
    def __init__( self ):
        ConsoleBase.__init__( self )

        self.savedspawn = os.spawnve

        sys.stdout = self
        sys.stderr = self

        win32.spawn = self.spawn
        posix.spawnvpe_spawn = self.spawn

    def spawn( self, sh, escape, cmd, args, env ): #IGNORE:W0613
        """
        Replaces SCons.Platform.spawn to colorize output of
        external processes
        """
        #import pdb;pdb.set_trace()
        #.replace( '"', '' )
        #FIXME: the " must be removed but not those with \ before them.
        tmp = []
        for x in args:
            x = x.replace(r'\"', '\???') # lets hope nobody wants to use \???
            x = x.replace(r'"', '')
            x = x.replace(r'\???', '"')
            tmp.append( subsitute_env_vars( x, env ) )
        args = tmp

        startupinfo = None
        if os.name == "nt":
            # Long command support <= 32766 for Windows
            startupinfo = sp.STARTUPINFO()
            startupinfo.dwFlags |= sp.STARTF_USESHOWWINDOW

        stdout = sp.PIPE
        #import pdb;pdb.set_trace()
        if ( os.name == "posix" and args[0] == "wine" ):
            stdout = None

        # We get unicode objects in the environment from somewhere, which makes
        # Popen unhappy. Force the environment to strings.
        # TODO(mika.raento): fix the source of the unicode.
        env = dict([ (k, str(v)) for (k, v) in env.iteritems() ])

        p = sp.Popen( args, bufsize = 1024,
                    stdout = stdout, stderr = sp.STDOUT,
                    startupinfo = startupinfo,
                    shell = False, env = env )
        result = None
        #import pdb;pdb.set_trace()
        if p.stdout is not None:
            while result is None:
                # This is slow on Linux with Wine!!
                line = p.stdout.read()
                self.write( line )
                result = p.poll()
            # Get the last lines
            line = p.stdout.readline()
            self.write( line )

        else:
            result = p.wait()

        return result

del __x


import torch
import sys

from vision.nn.mobilenet import MobileNetV1
from extract_tf_weights import read_weights

def fill_weights_torch_model(weights, state_dict):
    for name in state_dict:
        if name == 'classifier.weight':
            weight = weights['MobilenetV1/Logits/Conv2d_1c_1x1/weights']
            weight = torch.tensor(weight, dtype=torch.float32).permute(3, 2, 0, 1)
            assert state_dict[name].size() == weight.size()
            state_dict[name] = weight
        elif name == 'classifier.bias':
            bias = weights['MobilenetV1/Logits/Conv2d_1c_1x1/biases']
            bias = torch.tensor(bias, dtype=torch.float32)
            assert state_dict[name].size() == bias.size()
            state_dict[name] = bias
        elif name.endswith('BatchNorm.weight'):
            key = name.replace("features", "MobilenetV1").replace(".", "/").replace('BatchNorm/weight', 'BatchNorm/gamma')
            weight = torch.tensor(weights[key], dtype=torch.float32)
            assert weight.size() == state_dict[name].size()
            state_dict[name] = weight
        elif name.endswith('BatchNorm.bias'):
            key = name.replace("features", "MobilenetV1").replace(".", "/").replace('BatchNorm/bias', 'BatchNorm/beta')
            bias = torch.tensor(weights[key], dtype=torch.float32)
            assert bias.size() == state_dict[name].size()
            state_dict[name] = bias
        elif name.endswith('running_mean'):
            key = name.replace("features", "MobilenetV1").replace(".", "/").replace('running_mean', 'moving_mean')
            running_mean = torch.tensor(weights[key], dtype=torch.float32)
            assert running_mean.size() == state_dict[name].size()
            state_dict[name] = running_mean
        elif name.endswith('running_var'):
            key = name.replace("features", "MobilenetV1").replace(".", "/").replace('running_var', 'moving_variance')
            running_var = torch.tensor(weights[key], dtype=torch.float32)
            assert running_var.size() == state_dict[name].size()
            state_dict[name] = running_var
        elif name.endswith('depthwise.weight'):
            key = name.replace("features", "MobilenetV1").replace(".", "/")
            key = key.replace('depthwise/weight', 'depthwise/depthwise_weights')
            weight = torch.tensor(weights[key], dtype=torch.float32).permute(2, 3, 0, 1)
            assert weight.size() == state_dict[name].size()
            state_dict[name] = weight
        else:
            key = name.replace("features", "MobilenetV1").replace(".", "/").replace('weight', 'weights')
            weight = torch.tensor(weights[key], dtype=torch.float32).permute(3, 2, 0, 1)
            assert weight.size() == state_dict[name].size()
            state_dict[name] = weight

if __name__ == '__main__':
    if len(sys.argv) < 3:
        print("Usage: python translate_tf_modelnetv1.py <tf_model.pb> <pytorch_weights.pth>")
    tf_model = sys.argv[1]
    torch_weights_path = sys.argv[2]
    print("Extract weights from tf model.")
    weights = read_weights(tf_model)

    net = MobileNetV1(1001)
    states = net.state_dict()
    print("Translate tf weights.")
    fill_weights_torch_model(weights, states)
    torch.save(states, torch_weights_path)


# -*- coding: utf-8 -*-
'''
Manage Kibana objects.

.. code-block:: yaml

    kibana:
      kibana_url: 'https://es.host.com:9200'
      kibana_index: '.kibana'

.. code-block:: yaml

    Ensure minimum dashboard is managed:
      kibana_objects.present:
        - name: 'Logs'
        - kibana_content: <JSON object>
        - kibana_type: 'dashboard'

'''

# Import Python libs
import requests

# Import Salt libs
from salt.utils.dictdiffer import DictDiffer

def __virtual__():
    '''Always load the module.'''
    return True

def present(name, kibana_content=None, kibana_type=None):
    '''
    Ensure the Kibana object exists in the database.

    name
        Name of the object

    kibana_content
        Content in JSON

    kibana_type
        String
    '''
    ret = {'name': name, 'result': True, 'comment': '', 'changes': {}}

    if not kibana_content:
        ret['result'] = False
        ret['comment'] = 'Content is not set'
        return ret

    url, index = _get_parameters(name, kibana_type)
    if not url:
        ret['result'] = False
        ret['comment'] = index
        return ret

    try:
        headers = {'Content-type': 'application/json'}
        response = requests.put(url, headers=headers, json=kibana_content)
    except requests.exceptions.RequestException as exc:
        ret['result'] = False
        ret['comment'] = ("Failed to create Kibana object {0}\n"
                          "Got exception: {1}").format(name, exc)
    else:
        if response.ok:
            ret['comment'] = 'Kibana object {0} has been created'.format(name)
            ret['changes']['new'] = 'Kibana objects created'
        else:
            ret['result'] = False
            ret['comment'] = ("Failed to post Kibana object {0}\n"
                              "Response: {1}").format(name, response)

    return ret

def absent(name, kibana_type=None):
    '''
    Ensure the Kibana object is not present in the database.

    name
        Name of the object

    kibana_type
        String
    '''
    ret = {'name': name, 'result': True, 'comment': '', 'changes': {}}

    url, index = _get_parameters(name, kibana_type)
    if not url:
        ret['result'] = False
        ret['comment'] = index
        return ret

    try:
        response = requests.delete(url)
    except requests.exceptions.RequestException as exc:
        ret['result'] = False
        ret['comment'] = ("Failed to delete Kibana object {0}\n"
                          "Got exception: {1}").format(name, exc)
    else:
        if response.ok:
            ret['comment'] = "Kibana object {0} has been deleted".format(name)
        elif response.status_code == 404:
            ret['comment'] = "Kibana object {0} was not present".format(name)
        else:
            ret['result'] = False
            ret['comment'] = ("Failed to delete Kibana object {0}\n"
                              "Response: {1}").format(name, response)

    return ret

def _get_parameters(name, kibana_type):
    '''
    Retrieve parameters from profile.
    '''

    if not kibana_type:
        return False, 'Type is not set'

    profile = __salt__['config.option']('kibana')

    url = profile.get('kibana_url')
    if not url:
        return False, 'Cannot get URL needed by Kibana client'

    index = profile.get('kibana_index')
    if not index:
        return False, 'Cannot get the index needed by Kibana client'

    url = "http://{0}/{1}/{2}/{3}".format(url, index, kibana_type, name)
    return url, index


#!/usr/bin/env python3
import os
import sys
sys.path.append("../")
import utils
import imu
import time
from math import sqrt
import threading

stop_thread = threading.Event()
acc_data = []
imu_thread = threading.Thread(target=imu.imu_thread_func,args=(acc_data,stop_thread,))
imu_thread.start()

try:
    while True:
        time.sleep(0.075)
        if len(acc_data) > 0:
            os.system("clear")
            try:
                ax = round(acc_data[-1][0],3)
                ay = round(acc_data[-1][1],3)
                az = round(acc_data[-1][2],3)
            except:
                utils.brint("Missed data point",color="BOLD_RED")
            print(ax,ay,az)

except KeyboardInterrupt:
    imu.write_to_file(acc_data)
    stop_thread.set()


"""Dataset utilities."""

from typing import Union

import hashlib
import urllib.request

from tqdm import tqdm

def download_file(url: str, destination: str, progress_bar: bool = True, md5sum: Union[str, None] = None) -> None:
    """Download a file.

    Parameters
    ----------
    url: str
      File URL.
    destination: str
      Path to save the file.
    progress_bar: bool = True
      Show progress bar if True.
    md5sum: Union[str, None] = None
      md5sum hash of the file.

    Returns
    -------
    None
    """
    response = urllib.request.urlopen(url)
    file_size = int(response.info()["Content-Length"])

    def __show_progress(block_num: int, block_size: int, total_size: int) -> None:
        downloaded = block_num * block_size
        if total_size > 0:
            progress_bar.update(downloaded - progress_bar.n)

    with tqdm(total=file_size, unit='B', unit_scale=True, desc=destination, ncols=100) as progress_bar:
        urllib.request.urlretrieve(url, destination, __show_progress)

    if md5sum is not None:
        md5_hash = hashlib.md5()
        with open(destination, 'rb') as f:
            for chunk in iter(lambda: f.read(4096), b''):
                md5_hash.update(chunk)
        md5sum_test = md5_hash.hexdigest()
        if md5sum != md5sum_test:
            raise ValueError(f'md5sum mismatch. \nExpected: {md5sum}\nActual: {md5sum_test}')


from time import sleep
import argparse
import os
import re
import ntpath
import glob
import sys
import fileinput
from shutil import copyfile
import json

def parse_args():
    """Uses argparse to enable user to customize script functionality"""
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('-in_dir', '--input_directory', help='path to directory containing input files')
    parser.add_argument('-in_file', '--input_file_path', help='path to input file')
    parser.add_argument('-out', '--output_directory', default='./', help='directory to which the output directory "/.HaplotypeCaller/" will be written to')
    parser.add_argument('-n', '--num_cores', default='1', help='slurm job submission option')
    parser.add_argument('-cn', default="1", help='number of cores for Cromwell jobs')
    parser.add_argument('-t', '--runtime', default='2-00:00:00', help='slurm job submission option')
    parser.add_argument('-ct', default="3000", help='cromwell run time; please specify as number of minutes')
    parser.add_argument('-p', '--queue', default='long', help='slurm job submission option')
    parser.add_argument('--mem_per_cpu', default='15G', help='slurm job submission option')
    parser.add_argument('-cm', default='5000', help='cromwell cpu memory per core')
    parser.add_argument('--mail_type', default='FAIL', help='slurm job submission option')
    parser.add_argument('--mail_user', default='victor_mao@hms.harvard.edu', help='slurm job submission option')
    parser.add_argument('-cromwell', '--cromwell_path', default='/n/shared_db/singularity/hmsrc-gatk/cromwell-43.jar', help='path to cromwell.jar file')
    parser.add_argument('-r', default='/home/mk446/BiO/Install/GATK-bundle/2.8/b37/human_g1k_v37_decoy.fasta')
    parser.add_argument('-gatk', default='/home/mk446/BiO/Install/GATK4.1.2.0//gatk', help='path to software')
    parser.add_argument('-reference_name', default='b37', help='hg19, b37, etc.')
    parser.add_argument('-picard', default='/home/mk446/BiO/Install/picard-tools-2.5.0/picard.jar')
    parser.add_argument('-intervals', default='/n/data1/hms/dbmi/park/victor/references/Homo_sapiens_assembly19/interval_list.txt')

    return parser.parse_args()

def main():
    args = parse_args()
    clean_arg_paths(args)

    dirname = os.path.dirname(os.path.abspath(__file__))
    overrides = os.path.join(dirname, 'Overrides.config')
    wdl = os.path.join(dirname, 'Germline_workflow.wdl')
    json = os.path.join(dirname, 'Germline_workflow.json')

    input_files = return_input_files(args, 'bam') if args.input_directory is not None else [args.input_file_path]
    input_files = sort_by_size(input_files)

    for input_file in input_files:
        if not os.path.isfile(input_file): continue
        sample = os.path.basename(input_file).split('.')[0]
        sample_dir = os.path.join(args.output_directory,'.HaplotypeCaller/' + sample + '/')
        os.makedirs(sample_dir, exist_ok=True)

        input_json, input_config, input_wdl = generate_cromwell_inputs(args, input_file, json, wdl, overrides)
        slurm_command = return_slurm_command(args)
        output_file_name = gen_output_file_name(args, input_file)
        primary_command = return_primary_command(args, output_file_name, input_file, input_json, input_config, input_wdl)

        sh_file_name = gen_sh_file_name(args, output_file_name)
        write_out(args, slurm_command, primary_command, sh_file_name)

        sample_name = ntpath.basename(sh_file_name).split('.')[0] + '.vcf'
        sample_dir = os.path.join(ntpath.basename(sh_file_name).split('.')[0], sample_name)
        path_to_vcf = os.path.join(ntpath.dirname(ntpath.dirname(ntpath.dirname(sh_file_name))), sample_dir)
        #print(path_to_vcf)
        if not os.path.isfile(path_to_vcf):
            submit_job(sh_file_name)

def clean_arg_paths(args):
    """Modifies all user-inputted directory paths such that they end with a '/'"""
    d = vars(args)
    for arg in d.keys():
        if 'input_directory' in arg and d[arg]=='./': d[arg] = os.getcwd()
        if 'output_directory' in arg and d[arg]=='./': d[arg] = os.getcwd()
        if 'directory' in arg and d[arg] is not None and d[arg][-1] != '/': d[arg] += '/'

    output_dir = re.sub(" ", "", d["output_directory"])
    if output_dir[len(output_dir) - 1] is not "/":
        d["output_directory"] = output_dir + "/"

def return_input_files(args, ext):
    input_bams = [os.path.realpath(file) for file in glob.glob(args.input_directory + '*.' + ext)]
    return input_bams

def sort_by_size(input_files):
    for i in range(len(input_files)):
        input_files[i] = (input_files[i], os.path.getsize(input_files[i]))
    input_files.sort(key=lambda filename: filename[1])
    for i in range(len(input_files)):
        input_files[i] = input_files[i][0]
    return input_files

def generate_cromwell_inputs(args, input_file, json_file, wdl, overrides):
    dir = args.output_directory + '.HaplotypeCaller/' + '.' + os.path.basename(input_file).split('.')[0] + '/'
    os.makedirs(dir, exist_ok=True)

    bam_dir = os.path.dirname(input_file)
    bam_sample = os.path.basename(input_file)

    bai_suffix = '.bai'
    path = os.path.join(bam_dir, re.sub('.bam', '.bam.bai', bam_sample))
    if not (os.path.isfile(path) and os.access(path, os.R_OK)):
        path = os.path.join(bam_dir, re.sub('.bam', '.bai', bam_sample))

    copyfile(json_file, dir + 'Input.json')

    dict_path = os.path.dirname(args.r)
    ref = os.path.basename(args.r).split('.fa')[0]

    with open(dir + 'Input.json') as f:
        data = f.read()
        d = json.loads(data)
        d["HaplotypeCallerGvcf_GATK4.input_bam"] = input_file
        d["HaplotypeCallerGvcf_GATK4.input_bam_index"] = path
        d["HaplotypeCallerGvcf_GATK4.output_directory"] = os.path.join(args.output_directory,'.HaplotypeCaller/' + bam_sample.split('.')[0] + '/')
        d["HaplotypeCallerGvcf_GATK4.ref_dict"] = os.path.join(dict_path, ref + '.dict')
        d["HaplotypeCallerGvcf_GATK4.ref_fasta"] = args.r
        d["HaplotypeCallerGvcf_GATK4.ref_fasta_index"] = args.r + '.fai'
        d["HaplotypeCallerGvcf_GATK4.gatk_path"] = args.gatk
        d["HaplotypeCallerGvcf_GATK4.picard_path"] = args.picard
        d["HaplotypeCallerGvcf_GATK4.bam_directory"] = os.path.dirname(input_file)
        d["HaplotypeCallerGvcf_GATK4.scattered_calling_intervals_list"] = args.intervals

        if args.reference_name.lower() is not "b37":
            d["HaplotypeCallerGvcf_GATK4.reference"] = args.reference_name

    with open(dir + 'Input.json', 'w') as f:
        f.write(json.dumps(d))

    copyfile(overrides, dir + 'Overrides.config')

    with fileinput.FileInput(dir + 'Overrides.config', inplace=True) as file:
        for line in file:
            print(line.replace(
            "medium", args.queue).replace("!@#$", args.ct).replace("%^&*", args.cm).replace("Int cpus = 1", "Int cpus = " + args.cn), end='')

    if args.queue == 'park' or args.queue == 'priopark':
        f = open(dir + 'Overrides.config', "r")
        contents = f.readlines()
        f.close()

        contents.insert(95, "            --account=${account_name} \\\n")

        f = open(dir + 'Overrides.config', "w")
        contents = "".join(contents)
        f.write(contents)
        f.close()

    copyfile(wdl, os.path.join(dir, 'workflow.wdl'))

    return os.path.join(dir,'Input.json'), os.path.join(dir, 'Overrides.config'), os.path.join(dir, 'workflow.wdl')

def return_slurm_command(args):
    """Returns slurm command given args provided"""
    slurm_command = '#!/bin/bash\n' + \
                '#SBATCH -n ' + args.num_cores + '\n' + \
                '#SBATCH -t ' + args.runtime + '\n' + \
                '#SBATCH -p ' + args.queue + '\n' + \
                '#SBATCH --mem-per-cpu=' + args.mem_per_cpu + '\n' + \
                '#SBATCH --mail-type=' + args.mail_type + '\n' + \
                '#SBATCH --mail-user=' + args.mail_user + '\n' + \
        '#SBATCH --exclude=compute-p-17-[34-46]' + '\n'
    if args.queue in ['park', 'priopark']:
        slurm_command += '#SBATCH --account=park_contrib' + '\n'
    slurm_command += 'module load gcc/6.2.0 java/jdk-1.8u112 bcftools/1.9' + '\n'
    return slurm_command

def gen_output_file_name(args, input_file):
    sample = os.path.basename(input_file).split('.')[0]
    output_file_name = args.output_directory + '.HaplotypeCaller/' + '.' + sample + '/' + sample
    return output_file_name

def return_primary_command(args, output_file_name, input_file, input_json, input_config, input_wdl):
    primary_command = 'java -Dconfig.file=' + input_config + ' -jar ' + args.cromwell_path + ' run ' + input_wdl + ' -i ' + input_json
    return primary_command

def gen_sh_file_name(args, output_file_name):
    """Generates sh file name"""
    sh_file_name = os.path.dirname(output_file_name) + '/.sh/' + os.path.basename(output_file_name) + '.sh'
    return sh_file_name

def write_out(args, slurm_command, primary_command, sh_file_name):
    """"""
    os.makedirs(os.path.dirname(sh_file_name), exist_ok=True)
    with open(sh_file_name, 'w') as file:
        file.write(slurm_command + primary_command)

def submit_job(sh_file_name):
    os.chdir(os.path.dirname(sh_file_name))
    os.system('chmod +x ' + os.path.basename(sh_file_name))
    os.system('sbatch ./' + os.path.basename(sh_file_name))

if __name__ == "__main__":
    main()


import os
import sys
# from typing import Sequence
sys.path.insert(0,os.getcwd())
import copy
import argparse
import shutil
import time
import numpy as np
import random

import torch
# import torch.backends.cudnn as cudnn
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.nn.parallel import DataParallel

from utils.history import History
from utils.dataloader import Mydataset, collate
from utils.train_utils import train, validation, print_info, file2dict, init_random_seed, set_random_seed, resume_model
from utils.inference import init_model
from core.optimizers import *
from models.build import BuildNet

def parse_args():
    parser = argparse.ArgumentParser(description='Train a model')
    parser.add_argument('config', help='train config file path')
    parser.add_argument('--resume-from', help='the checkpoint file to resume from')
    parser.add_argument('--seed', type=int, default=None, help='random seed')
    parser.add_argument('--device', help='device used for training. (Deprecated)')
    parser.add_argument(
        '--gpu-id',
        type=int,
        default=0,
        help='id of gpu to use '
        '(only applicable to non-distributed training)')
    parser.add_argument(
        '--split-validation',
        action='store_true',
        help='whether to split validation set from training set.')
    parser.add_argument(
        '--ratio',
        type=float,
        default=0.2,
        help='the proportion of the validation set to the training set.')
    parser.add_argument(
        '--deterministic',
        action='store_true',
        help='whether to set deterministic options for CUDNN backend.')
    parser.add_argument('--local-rank', type=int, default=0)
    args = parser.parse_args()
    if 'LOCAL_RANK' not in os.environ:
        os.environ['LOCAL_RANK'] = str(args.local_rank)
    return args

def main():
    # 读取配置文件获取关键字段
    args = parse_args()
    model_cfg,train_pipeline,val_pipeline,data_cfg,lr_config,optimizer_cfg = file2dict(args.config)
    print_info(model_cfg)

    # 初始化
    meta = dict()
    dirname = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
    save_dir = os.path.join('logs',model_cfg.get('backbone').get('type'),dirname)
    meta['save_dir'] = save_dir

    # 设置随机数种子
    seed = init_random_seed(args.seed)
    set_random_seed(seed, deterministic=args.deterministic)
    meta['seed'] = seed

    # 读取训练&制作验证标签数据
    total_annotations   = "datas/train.txt"
    with open(total_annotations, encoding='utf-8') as f:
        total_datas = f.readlines()
    if args.split_validation:
        total_nums = len(total_datas)
        # indices = list(range(total_nums))
        if isinstance(seed, int):
            rng = np.random.default_rng(seed)
            rng.shuffle(total_datas)
        val_nums = int(total_nums * args.ratio)
        folds = list(range(int(1.0 / args.ratio)))
        fold = random.choice(folds)
        val_start = val_nums * fold
        val_end = val_nums * (fold + 1)
        train_datas = total_datas[:val_start] + total_datas[val_end:]
        val_datas = total_datas[val_start:val_end]
    else:
        train_datas = total_datas.copy()
        test_annotations    = 'datas/test.txt'
        with open(test_annotations, encoding='utf-8') as f:
            val_datas   = f.readlines()

    # 初始化模型,详见https://www.bilibili.com/video/BV12a411772h
    if args.device is not None:
        device = torch.device(args.device)
    else:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print('Initialize the weights.')
    model = BuildNet(model_cfg)
    if not data_cfg.get('train').get('pretrained_flag'):
        model.init_weights()
    if data_cfg.get('train').get('freeze_flag') and data_cfg.get('train').get('freeze_layers'):
        freeze_layers = ' '.join(list(data_cfg.get('train').get('freeze_layers')))
        print('Freeze layers : ' + freeze_layers)
        model.freeze_layers(data_cfg.get('train').get('freeze_layers'))

    if device != torch.device('cpu'):
        model = DataParallel(model,device_ids=[args.gpu_id])

    # 初始化优化器
    optimizer = eval('optim.' + optimizer_cfg.pop('type'))(params=model.parameters(),**optimizer_cfg)

    # 初始化学习率更新策略
    lr_update_func = eval(lr_config.pop('type'))(**lr_config)

    # 制作数据集->数据增强&预处理,详见https://www.bilibili.com/video/BV1zY4y167Ju
    train_dataset = Mydataset(train_datas, train_pipeline)
    val_pipeline = copy.deepcopy(train_pipeline)
    val_dataset = Mydataset(val_datas, val_pipeline)
    train_loader = DataLoader(train_dataset, shuffle=True, batch_size=data_cfg.get('batch_size'), num_workers=data_cfg.get('num_workers'),pin_memory=True, drop_last=True, collate_fn=collate)
    val_loader = DataLoader(val_dataset, shuffle=False, batch_size=data_cfg.get('batch_size'), num_workers=data_cfg.get('num_workers'), pin_memory=True,
    drop_last=True, collate_fn=collate)

    # 将关键字段存储，方便训练时同步调用&更新
    runner = dict(
        optimizer         = optimizer,
        train_loader      = train_loader,
        val_loader        = val_loader,
        iter              = 0,
        epoch             = 0,
        max_epochs       = data_cfg.get('train').get('epoches'),
        max_iters         = data_cfg.get('train').get('epoches')*len(train_loader),
        best_train_loss   = float('INF'),
        best_val_acc     = float(0),
        best_train_weight = '',
        best_val_weight   = '',
        last_weight       = ''
    )
    meta['train_info'] = dict(train_loss = [],
                              val_loss = [],
                              train_acc = [],
                              val_acc = [])

    # 是否从中断处恢复训练
    if args.resume_from:
        model,runner,meta = resume_model(model,runner,args.resume_from,meta)
    else:
        os.makedirs(save_dir)
        shutil.copyfile(args.config,os.path.join(save_dir,os.path.split(args.config)[1]))
        model = init_model(model, data_cfg, device=device, mode='train')

    # 初始化保存训练信息类
    train_history =History(meta['save_dir'])

    # 记录初始学习率，详见https://www.bilibili.com/video/BV1WT4y1q7qN
    lr_update_func.before_run(runner)

    # 训练
    for epoch in range(runner.get('epoch'),runner.get('max_epochs')):
        lr_update_func.before_train_epoch(runner)
        train(model,runner, lr_update_func, device, epoch, data_cfg.get('train').get('epoches'), meta)
        validation(model,runner, data_cfg.get('test'), device, epoch, data_cfg.get('train').get('epoches'), meta)

        train_history.after_epoch(meta)

if __name__ == "__main__":
    main()


import sqlite3
from datetime import datetime

conn = sqlite3.connect('crud.db')
cursor = conn.cursor()

def main():
        print ("Qual tabela gostaria de acessar? ")
        print ("Para acessar postagens, escreva POST" )
        print ("Para acessar comentarios, escreva COMMENT ")
        table = input("-> ")

        if table.lower() == "post" or table.lower() == "comment":
                crud(table)
        else:
                print("Informe um nome válido.")

def crud(table):
        x = input("Qual ação deseja realizar: Create(C), Read(R), Update(U) ou Delete(D)?")
        xLower = x.lower()
        if table.lower() == "post":
                if xLower == "c":
                        c_post()
                elif xLower == "r":
                        r_post()
                elif xLower == "u":
                        u_post()
                elif xLower == "d":
                        d_post()
                else:
                        print("Insira uma ação válida!!")
        else:
                if xLower == "c":
                        c_comment()
                elif xLower == "r":
                        r_comment()
                elif xLower == "u":
                        u_comment()
                elif xLower == "d":
                        d_comment()
                else:
                        print("Insira uma ação válida!!")

def c_post():
        title = input("Qual o titulo da postagem? ")
        created = datetime.now()
        text = input("Qual vai ser o conteúdo? ")
        cursor.execute("INSERT INTO post (Title, Created, Text) VALUES (?, ?, ?)", (title, created, text))
        conn.commit()
        print("Sucesso :)")

def c_comment():
        postIdAdd = input("Insira o ID da postagem ")
        textAdd = input("Qual será o comentário? ")
        created = datetime.now()
        userAdd = input("Qual é o nome do usuario? ")
        cursor.execute("""
        SELECT PostID FROM post WHERE PostID == ?;
        """, postIdAdd)
        postId = cursor.fetchone()

        if postId == None:
                print("Algo deu errado :( ")

        else:
                cursor.execute("INSERT INTO comment (PostID, Text, Created, User) VALUES (?, ?, ?, ?)", (postId[0], textAdd, created, userAdd))

                conn.commit()

                print("Sucesso :)")

def r_post():
        cursor.execute("""
        SELECT * FROM post;
        """)

        for line in cursor.fetchall():
                print(line)

def r_comment():
        cursor.execute("""
        SELECT * FROM comment;
        """)

        for line in cursor.fetchall():
                print(line)

def u_post():
        update = input("Qual é o ID da postagem que quer alterar? ")

        cursor.execute("""
        SELECT PostID FROM post
        WHERE PostID == ?
        """, update)
        checkPost = cursor.fetchone()

        if checkPost == None:
                print("Não foi possivel encontrar esse ID ")

        else:
                data = input("Informe o que deseja alterar (titulo, texto) ")

                if data.lower() == "titulo":
                        newData = input("Insira o novo titulo ")

                        cursor.execute("""
                        UPDATE post SET Title = ?
                        WHERE PostID = ?
                        """, (newData, update))

                        conn.commit()

                        print("Sucesso :) ")

                elif data.lower() == "texto":
                        newData = input("Insira qual é a edição: ")

                        cursor.execute("""
                        UPDATE post SET Text = ?
                        WHERE PostID
                        """, (newData, update))

                        conn.commit()

                        print("Sucesso :) ")

                else:
                        print("Invalido")

def u_comment():
        update = input("Qual é o ID do comentario que quer alterar? ")

        cursor.execute("""
        SELECT CommentID FROM comment
        WHERE CommentID == ?
        """, update)
        checkComment = cursor.fetchone()

        if checkComment == None:
                print("Não foi possivel encontrar esse ID ")

        else:
                data = input("Informe o que deseja alterar (Texto, usuario) ")

                if data.lower() == "text":
                        newData = input("Insira o novo texto ")

                        cursor.execute("""
                        UPDATE comment SET Text = ?
                        WHERE CommentID = ?
                        """, (newData, update))

                        conn.commit()

                        print("Sucesso :) ")

                elif data.lower() == "usuario":
                        newData = input("Insira qual é o usuario: ")

                        cursor.execute("""
                        UPDATE comment SET User = ?
                        WHERE CommentID = ?
                        """, (newData, update))

                        conn.commit()

                        print("Sucesso :) ")

                else:
                        print("Invalido")

def d_post():
        delete = input("Qual é o ID do post a ser excluido? ")

        cursor.execute("""
        SELECT PostID from post
        WHERE PostID == ?
        """, delete)
        checkPost = cursor.fetchone()

        if checkPost == None:
                print("ID invalido ")
        else:
                cursor.execute("""
                DELETE FROM post
                WHERE PostID == ?
                """, delete)

                cursor.execute("""
                DELETE FROM comment
                WHERE PostID == ?
                """, delete)

                conn.commit()

                print("A postagem foi deletada com sucesso :)")

def d_comment():
        delete = input("Qual é o ID do comentario? ")

        cursor.execute("""
        SELECT CommentID FROM comment
        WHERE CommentID == ?
        """, delete)
        checkComment = cursor.fetchone()

        if checkComment == None:
                print("Não foi possivel encontrar esse comentario ")
        else:
                cursor.execute("""
                DELETE FROM comment
                WHERE CommentID == ?
                """, delete)

                conn.commit()

                print("Sucesso :)")

main()


import pandas as pd
import simplejson as json

data_file = open("fiveyearsdata.json", "r")
data = json.load(data_file)
data = data

print("Data size:", len(data), "observations")
df = pd.DataFrame()
for i, d in enumerate(data[:]): # Going through data: observation by observation per each day
    try:
        df_1 = pd.DataFrame(d["data"]["weather"], index=[i])
        df_2 = pd.DataFrame(d["data"]["request"], index=[i])
        df_ = pd.concat([df_1, df_2], axis=1)
        for k, v in d["data"]["weather"][0]["hourly"][0].items():
            df_.loc[i, k] = v
        for k, v in d["data"]["weather"][0]["astronomy"][0].items():
            df_.loc[i, k] = v

        df = df_ if df.empty else df.append(df_)
    except:
        # There are some errors found in fetched data already.
        # Days missed as a result of an error should be imputed by some way
        print("Error at index:", i, "Error details:", d["data"])

df = df.drop('hourly', 1)
df = df.drop('astronomy', 1)

print("Fetched data sieze:", len(df), "observations")
print("Number of variables:", len(df.columns))

df.to_csv('weather_data.csv', sep='|') # sep here is to avoid "," in the query column


# encoding=<Encoding.LATIN1: 0>

import os
from shutil import copy2, move
from pydub import AudioSegment

import mutagen
from mutagen import mp4, easymp4, mp3, aiff

class AifMetaData:
  def __init__(self, fl):
    self.title = str(fl.get('TIT2', None))
    self.artist = str(fl.get('TPE1', None))
    self.albumArtist = str(fl.get('TPE2', None))
    self.composer = str(fl.get('TCOM', None))
    self.album = str(fl.get('TALB', None))
    self.trackNo = str(fl.get('TRCK', None)[0]).split('/')[0]
    self.tracks = str(fl.get('TRCK', None)[0]).split('/')[1]
    self.diskNo = str(fl.get('TPOS', None)[0]).split('/')[0]
    self.disks = str(fl.get('TPOS', None)[0]).split('/')[1]
    self.genre = str(fl.get('TCON', None))
    self.year = str(fl.get('TDRC', None))

  def summary(self):
    varsCopy = vars(self).copy()
    return varsCopy

walk_path = os.path.join('G:', 'My Drive', 'Music', 'aif by Album')

for source_dir, _, files in os.walk(walk_path):

    try:
      print(source_dir)
    except:
      continue

    for file in files:

        fileName = file.split('.')[:-1]
        fileExt = file.split('.')[-1]
        print(fileName, fileExt)
        quit()

        if not file.lower().endswith(('aif')):
          continue

        source_path = os.path.join(source_dir, file)

        mut = aiff.AIFF(source_path)
        metaData = AifMetaData(mut)

        aif = AudioSegment.from_file(source_path)


#!/usr/bin/env python
# lint_ignore=E501
# ***** BEGIN LICENSE BLOCK *****
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this file,
# You can obtain one at http://mozilla.org/MPL/2.0/.
# ***** END LICENSE BLOCK *****
""" updates.py

A script publish a release to Balrog.

"""

import os
import sys
from datetime import datetime, timedelta

sys.path.insert(1, os.path.dirname(os.path.dirname(sys.path[0])))
from mozharness.base.vcs.vcsbase import MercurialScript
from mozharness.mozilla.buildbot import BuildbotMixin
from mozharness.base.log import FATAL

# PublishBalrog {{{1

class PublishBalrog(MercurialScript, BuildbotMixin):

    def __init__(self, require_config_file=True):
        super(PublishBalrog, self).__init__(
            all_actions=[
                'clobber',
                'pull',
                'submit-to-balrog',
            ],
            default_actions=[
                'clobber',
                'pull',
                'submit-to-balrog',
            ],
            config={
                'buildbot_json_path': 'buildprops.json',
                'credentials_file': 'oauth.txt',
            },
            require_config_file=require_config_file
        )

    def _pre_config_lock(self, rw_config):
        super(PublishBalrog, self)._pre_config_lock(rw_config)
        # override properties from buildbot properties here as defined by
        # taskcluster properties
        self.read_buildbot_config()
        if not self.buildbot_config:
            self.warning("Skipping buildbot properties overrides")
            return
        # TODO: version and appVersion should come from repo
        props = self.buildbot_config["properties"]
        for prop in ['product', 'version', 'build_number', 'channels',
                     'balrog_api_root', 'schedule_at', 'background_rate',
                     'publish_bz2_blob']:
            if props.get(prop):
                self.info("Overriding %s with %s" % (prop, props[prop]))
                self.config[prop] = props.get(prop)

    def query_abs_dirs(self):
        if self.abs_dirs:
            return self.abs_dirs
        self.abs_dirs = super(PublishBalrog, self).query_abs_dirs()
        self.abs_dirs["abs_tools_dir"] = os.path.join(
            self.abs_dirs['abs_work_dir'], self.config["repo"]["dest"])
        return self.abs_dirs

    def query_channel_configs(self):
        """Return a list of channel configs.
        For RC builds it returns "release" and "beta" using
        "enabled_if_version_matches" to match RC.

        :return: list
         """
        return [(n, c) for n, c in self.config["update_channels"].items() if
                n in self.config["channels"]]

    def query_repos(self):
        """Build a list of repos to clone."""
        return [self.config["repo"]]

    def pull(self):
        super(PublishBalrog, self).pull(
            repos=self.query_repos())

    def submit_to_balrog(self):
        for _, channel_config in self.query_channel_configs():
            self._submit_to_balrog(channel_config)
        if 'publish_bz2_blob' in self.config and \
                self.config['publish_bz2_blob']:
            for _, channel_config in self.query_channel_configs():
                self._submit_to_balrog_bz2(channel_config)

    def _submit_to_balrog(self, channel_config):
        dirs = self.query_abs_dirs()
        auth = os.path.join(os.getcwd(), self.config['credentials_file'])
        cmd = [
            sys.executable,
            os.path.join(dirs["abs_tools_dir"],
                         "scripts/build-promotion/balrog-release-shipper.py")]
        cmd.extend([
            "--api-root", self.config["balrog_api_root"],
            "--credentials-file", auth,
            "--username", self.config["balrog_username"],
            "--version", self.config["version"],
            "--product", self.config["product"],
            "--build-number", str(self.config["build_number"]),
            "--verbose",
        ])
        for r in channel_config["publish_rules"]:
            cmd.extend(["--rules", str(r)])
        if channel_config.get("schedule_asap"):
            # RC releases going to the beta channel have no ETA set for the
            # RC-to-beta push. The corresponding task is scheduled after we
            # resolve the push-to-beta human decision task, so we can schedule
            # it ASAP plus some additional 30m to avoid retry() to fail.
            schedule_at = datetime.utcnow() + timedelta(minutes=30)
            cmd.extend(["--schedule-at", schedule_at.isoformat()])
        elif self.config.get("schedule_at"):
            cmd.extend(["--schedule-at", self.config["schedule_at"]])
        if self.config.get("background_rate"):
            cmd.extend(["--background-rate", str(self.config["background_rate"])])

        self.retry(lambda: self.run_command(cmd, halt_on_failure=True),
                   error_level=FATAL)

    def _submit_to_balrog_bz2(self, channel_config):
        dirs = self.query_abs_dirs()
        # Use env varialbe instead of command line to avoid issues with blob
        # names starting with "-", e.g. "-bz2"
        env = {"BALROG_BLOB_SUFFIX": channel_config["bz2_blob_suffix"]}
        auth = os.path.join(os.getcwd(), self.config['credentials_file'])
        cmd = [
            sys.executable,
            os.path.join(dirs["abs_tools_dir"],
                         "scripts/build-promotion/balrog-release-shipper.py")]
        cmd.extend([
            "--api-root", self.config["balrog_api_root"],
            "--credentials-file", auth,
            "--username", self.config["balrog_username"],
            "--version", self.config["version"],
            "--product", self.config["product"],
            "--build-number", str(self.config["build_number"]),
            "--suffix", channel_config["bz2_blob_suffix"],
            "--verbose",
        ])
        for r in channel_config["bz2_publish_rules"]:
            cmd.extend(["--rules", str(r)])
        if channel_config.get("schedule_asap"):
            # RC releases going to the beta channel have no ETA set for the
            # RC-to-beta push. The corresponding task is scheduled after we
            # resolve the push-to-beta human decision task, so we can schedule
            # it ASAP plus some additional 30m to avoid retry() to fail.
            schedule_at = datetime.utcnow() + timedelta(minutes=30)
            cmd.extend(["--schedule-at", schedule_at.isoformat()])
        elif self.config.get("schedule_at"):
            cmd.extend(["--schedule-at", self.config["schedule_at"]])
        if self.config.get("background_rate"):
            cmd.extend(["--background-rate", str(self.config["background_rate"])])

        self.retry(lambda: self.run_command(cmd, halt_on_failure=True, env=env),
                   error_level=FATAL)

# __main__ {{{1
if __name__ == '__main__':
    PublishBalrog().run_and_exit()


from selenium.webdriver.common.by import By
from traceback import print_stack
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import *
import utilities.custom_logger as cl
import logging
import time
import os

class SeleniumDriver():

    log = cl.customLogger(logging.DEBUG)

    def __init__(self, driver):
        self.driver = driver

    def getByType(self, locatorType):
        """
        returns Bytype
        """

        locatorType = locatorType.lower()

        if locatorType == "id":
            return By.ID
        elif locatorType == "name":
            return By.NAME
        elif locatorType == "xpath":
            return By.XPATH
        elif locatorType == "css":
            return By.CSS_SELECTOR
        elif locatorType == "class":
            return By.CLASS_NAME
        elif locatorType == "link":
            return By.LINK_TEXT
        else:
            self.log.info("Locator type " + locatorType +
                          " is not correct")
        return False

    def getElement(self, locator, locatorType="id"):

        element = None
        try:
            locatorType = locatorType.lower()
            byType = self.getByType(locatorType)
            element = self.driver.find_element(byType, locator)
            self.log.info("Succesfully found element with locator: " + locator +
                          " and  locatorType: " + locatorType)
        except:
            self.log.info("Unable to find Element with locator: " + locator +
                          " and  locatorType: " + locatorType)
        return element

    def elementClick(self, locator, locatorType="id"):

        try:
            element = self.getElement(locator, locatorType)
            element.click()
            self.log.info("Successfully clicked the element with locator: " + locator +
                          " locatorType: " + locatorType)
        except:
            self.log.info("Unable to click on the element with locator: " + locator +
                          " locatorType: " + locatorType)
            print_stack()

    def sendKeys(self, data, locator, locatorType="id"):
        try:
            element = self.getElement(locator, locatorType)
            element.send_keys(data)
            self.log.info("Succssfully Sent data with locator: " + locator +
                          " locatorType: " + locatorType)
        except:
            self.log.info("Unable to send data with locator: " + locator +
                  " locatorType: " + locatorType)
            print_stack()

    def waitForElement(self, locator, locatorType="id",
                               timeout=10, pollFrequency=0.5):
        element = None
        try:
            byType = self.getByType(locatorType)

            wait = WebDriverWait(self.driver, 10, poll_frequency=1,
                                 ignored_exceptions=[NoSuchElementException,
                                                     ElementNotVisibleException,
                                                     ElementNotSelectableException])
            element = wait.until(EC.element_to_be_clickable((byType,locator)))

            self.log.info("Successfully Element appeared on the web page")
        except:
            self.log.info("Unable to find Element on the web page")
            print_stack()
        return element

    def webScroll(self, direction="up"):

        if direction == "up":
            self.driver.execute_script("window.scrollBy(0, -1000);")

        if direction == "down":
            self.driver.execute_script("window.scrollBy(0, 900);")

    def getText(self, locator="", locatorType="id"):

        try:
            element = self.getElement(locator, locatorType)
            text = element.text
            if len(text) != 0:
                text = text.strip()

        except:
            self.log.error("Failed to get text on element ")
            print_stack()
            text = None
        return text

    def screenShot(self, resultMessage):

        fileName = resultMessage + "." + str(round(time.time() * 1000)) + ".png"
        screenshotDirectory = "../screenshots/"
        relativeFileName = screenshotDirectory + fileName
        currentDirectory = os.path.dirname(__file__)
        destinationFile = os.path.join(currentDirectory, relativeFileName)
        destinationDirectory = os.path.join(currentDirectory, screenshotDirectory)

        try:
            if not os.path.exists(destinationDirectory):
                os.makedirs(destinationDirectory)
            self.driver.save_screenshot(destinationFile)
            self.log.info("Screenshot save to directory: " + destinationFile)
        except:
            self.log.error("### Exception Occurred when taking screenshot")
            print_stack()


from django.db.models.signals import post_save
from django.dispatch import receiver

from .models import User, LoanedBook
from .tasks import send_user_data, send_loaned_book_data

@receiver(post_save, sender=User)
def send_user_data_event(sender, **kwargs) -> None:
    user = kwargs["instance"]
    created = kwargs["created"]

    # Call the celery task to send user data to the admin_api service
    send_user_data.delay(user.id, user.email, user.first_name, user.last_name, created)

@receiver(post_save, sender=LoanedBook)
def send_loaned_book_data_event(sender, **kwargs) -> None:
    loaned_book: LoanedBook = kwargs["instance"]

    # Call the celery task to send the bood data that has been loaned to the admin_api service
    send_loaned_book_data.apply_async(
        kwargs={
            "date_borrowed": loaned_book.date_borrowed,
            "return_date": loaned_book.return_date,
            "book_id": loaned_book.book.id,
            "user_id": loaned_book.user.id,
        }
    )


import config as cfg

if cfg.lang == 'en':
    import localization.en as msg_cfg
else:
    import localization.ru as msg_cfg

from messages.message_base import MessageBase

class ModuleMsg(MessageBase):
    """Constructs the module asking message with Buttons
     based on MessageBase class"""

    def __init__(self, channel):
        MessageBase.__init__(self, channel)

    def __call__(self):
        return {
            "ts": self.timestamp,
            "channel": self.channel,
            "username": self.username,
            "icon_emoji": self.icon_emoji,
            "blocks": [
                self.make_section_block(msg_cfg.module_message['select']),
                self.get_buttons(),
            ],
        }

    @staticmethod
    def get_buttons():
        def button(text):
            return {
                "type": 'button',
                "text": {
                    "type": "plain_text",
                    "text": text,
                    "emoji": False,
                }
            }
        return {
            "type": 'actions',
            "elements": [button(msg_cfg.module_message['button'].format(i))
                         for i in range(1, cfg.module_number+1)],
        }


# -*- coding: utf-8 -*-
from __future__ import unicode_literals

from django.db import models, migrations
import datetime

class Migration(migrations.Migration):

    dependencies = [
        ('project', '0001_initial'),
        ('inventory', '0001_initial'),
    ]

    operations = [
        migrations.CreateModel(
            name='Demand',
            fields=[
                ('id', models.AutoField(primary_key=True, serialize=False, auto_created=True, verbose_name='ID')),
                ('purpose', models.CharField(max_length=254)),
                ('date', models.DateField(default=datetime.datetime.today)),
                ('site', models.ForeignKey(related_name='demands', to='project.Project')),
            ],
        ),
        migrations.CreateModel(
            name='DemandRow',
            fields=[
                ('id', models.AutoField(primary_key=True, serialize=False, auto_created=True, verbose_name='ID')),
                ('purpose', models.CharField(max_length=100, blank=True, null=True)),
                ('quantity', models.FloatField()),
                ('unit', models.CharField(max_length=50)),
                ('fulfilled_quantity', models.FloatField()),
                ('status', models.BooleanField(default=False)),
                ('demand', models.ForeignKey(related_name='rows', to='inventory.Demand')),
                ('item', models.ForeignKey(to='inventory.Item')),
            ],
        ),
    ]


import os
import logging
from collections import OrderedDict
import json
import numpy as np
import random

import torch

def set_logger(params, log_file=None):
    if log_file is None:
        dataset_id = params['dataset_id']
        model_id = params['model_id']
        log_dir = os.path.join(params['model_root'], dataset_id)
        log_file = os.path.join(log_dir, model_id + '.log')
    log_dir = os.path.dirname(log_file)
    os.makedirs(log_dir, exist_ok=True)

    # logs will not show in the file without the two lines.
    for handler in logging.root.handlers[:]:
        logging.root.removeHandler(handler)

    logging.basicConfig(level=logging.INFO,
                        format='%(asctime)s P%(process)d %(levelname)s %(message)s',
                        handlers=[logging.FileHandler(log_file, mode='w'),
                                  logging.StreamHandler()])

def print_to_json(data, sort_keys=True):
    new_data = dict((k, str(v)) for k, v in data.items())
    if sort_keys:
        new_data = OrderedDict(sorted(new_data.items(), key=lambda x: x[0]))
    return json.dumps(new_data, indent=4)

def seed_everything(seed=1029):
    random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True


import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import sklearn.metrics
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler
from sklearn.cluster import KMeans
import os

os.environ['OMP_NUM_THREADS'] = '2'

data = pd.read_csv('pytorch/data/sales data.csv')
print(data.head())

categorical_features = ['Channel', 'Region']
continuous_features = ['Fresh',  'Milk',  'Grocery',  'Frozen',  'Detergents_Paper',  'Delicassen']

for col in categorical_features:
    dummies = pd.get_dummies(data[col], prefix=col)
    data = pd.concat([data, dummies], axis=1)
    data.drop(col, axis=1, inplace=True)
print(data.head())

mms = MinMaxScaler()
mms.fit(data)
data_transformed = mms.transform(data)
print(data_transformed)

sum_of_squared_distances = []
K = range(1, 15)
for k in K:
    km = KMeans(n_clusters=k)
    km = km.fit(data_transformed)
    print(km.labels_)
    sum_of_squared_distances.append(km.inertia_)

# print(sum_of_squared_distances)
plt.plot(K, sum_of_squared_distances, 'bx-')
plt.xlabel('k')
plt.ylabel('sum_of_squared_distances')
plt.title('Optimal K')
plt.show()


import copy
import os
import time
import warnings
from collections import defaultdict, namedtuple
from enum import unique
from typing import Callable, Dict, Iterable, List, NamedTuple, Optional, Tuple, Union
from continuum.rehearsal.memory import RehearsalMemory
from continuum.tasks.base import BaseTaskSet

import numpy as np
import torch
from itertools import cycle
import torch.nn.functional as F
import torch.multiprocessing as mp
from fvcore.nn import FlopCountAnalysis
from simple_parsing import ArgumentParser, choice
from torch import nn
from tqdm import tqdm
from torch.utils.data import DataLoader, dataloader
import wandb
from continuum.tasks.task_set import TaskSet
from continuum.scenarios import _BaseScenario
from args import ArgsGenerator, ExperimentState
from continuum.tasks import TaskType, concat #, split_train_val
# from continuum.tasks.base import BaseTaskSet
from continuum.tasks.h5_task_set import H5TaskSet
from dataset_encoder import UsedFlops, estimate_compute_regime, prepare_scenarios
from Models import Classifier_options
from torchvision.transforms.transforms import Compose, ToTensor
from Models.model import (ModelContainer, ModelContainer_ER, TorchModuleWrapper)
from Utils.utils import SumMeter, is_connected, log_wandb, set_seed, split_er_buffer, get_scenario_remapping, preprocess_tensor, Logger, test, FlopsMeter,  prepare_dataloader, split_train_val, prepare_balanced_train_loader, prepare_randombuffer_train_loader, convert_to_task_set, TinyDB_hp_tracker
import logging

device = 'cuda' if torch.cuda.is_available() else 'cpu'
#used_flops_tuple=namedtuple("ResultTuple",['total_gigaflops_fw','total_gigaflops_bw','total_gigaflops','total_gigaflops_fw_es','total_gigaflops_bw_es', 'total_gigaflops_es','walltime_task','total_flops_fw_epoch', 'total_flops_bw_epoch', 'walltime_epoch'])
loss_function = nn.CrossEntropyLoss()

def init_flops_meters()-> Tuple[FlopsMeter,FlopsMeter]:
    flopws_meter_pertask_cv = FlopsMeter(['Gigaflops_sofar_fw(including_pt_cv)','Gigaflops_sofar_bw(including_pt_cv)','Gigaflops_sofar(including_pt_cv)',
                                        'Gigaflops_sofar_fw_erlst(including_pt_cv)', 'Gigaflops_sofar_bw_erlst(including_pt_cv)', 'Gigaflops_sofar_erlst(including_pt_cv)', 'walltime_task(including_pt_cv)'])
    flopws_meter_optimal_params = FlopsMeter(['Gigaflops_sofar_fw','Gigaflops_sofar_bw','Gigaflops_sofar',
                                        'Gigaflops_sofar_fw_erlst', 'Gigaflops_sofar_bw_erlst', 'Gigaflops_sofar_erlst', 'walltime_task'])
    return flopws_meter_pertask_cv, flopws_meter_optimal_params

def train_sklearn_classifier(model, task_id, train_loader, test_loaders, flops_meters:FlopsMeter):
    if isinstance(model.classifiers, Iterable):
            model_logreg = model.classifiers[task_id]
    else:
        model_logreg = model.classifiers
    #collect data
    x_train=[]
    y_train=[]
    x_test=[]
    y_test=[]
    for x,y,_ in train_loader:
        x,y=preprocess_tensor(model,x,y)
        x_train.append(x)
        y_train.append(y)
    for x,y,_ in test_loaders[-1]:
        x,y=preprocess_tensor(model,x,y)
        x_test.append(x)
        y_test.append(y)
    x_train,y_train=torch.cat(x_train).squeeze().cpu().numpy(),torch.cat(y_train).cpu().numpy()
    x_test,y_test=torch.cat(x_test).squeeze().cpu().numpy(),torch.cat(y_test).cpu().numpy()
    #this throws an error somehow
    # x_train, y_train, _ = train_loader.dataset.get_raw_samples()
    # x_test, y_test, _ = test_loaders[-1].get_raw_samples()
    model_logreg.fit(x_train,y_train)
    # Evaluate using the logistic regression classifier
    predictions = model_logreg.predict(x_test)
    accuracy_valid = np.mean((y_test == predictions).astype(float)) #* 100.
    print(f"Accuracy = {accuracy_valid:.3f}")
    # result=used_flops_tuple(*[v for v in flops_meters.value().values()],0,0)
    return model, accuracy_valid, None

def train_model(model_container:ModelContainer_ER,  args:ArgsGenerator, model_name:str, task_id:int, model:TorchModuleWrapper, train_loader:DataLoader, val_loader:DataLoader, train_loaders:List[DataLoader]=None, test_loaders:List=None, lr_anneal=None, epochs=None):
    #task level cross validation
    if lr_anneal is None:
        lr_anneal=model_container.args.lr_anneal
    best_valid_acc=0.
    best_model= None
    n_epochs_without_improvement = 0
    early_stopped = 0
    flops_meters = FlopsMeter(['flops_task_forward','flops_task_backward', 'flops_task', 'flops_task_forward_es', 'flops_task_backward_es', 'flops_task_es'])
    flops_per_epoch_fw=0
    flops_per_epoch_bw=0
    time_per_batch = None

    if model_container.args.classifier_type in ['logistic_regression', 'random_forrest', 'knn']:
        return train_sklearn_classifier(model, task_id, train_loader, test_loaders, flops_meters)
    else:
        if epochs is None:
            epochs = args.epochs
        convergence_epoch=epochs
        best_model_epoch=epochs
        flops_per_batch_fw=defaultdict(lambda: None)
        flops_per_batch_bw=defaultdict(lambda:None)
        lr_scheduler=None
        if model.optimizer is not None:
            for pg_i, param_group in enumerate(model.optimizer.param_groups):
                param_group["lr"] = model.lr
            lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(model.optimizer, T_max=args.epochs)

        if model.optimizer is None:
            epochs = 1

        len_train_loader=len(train_loader)
        _train_loader=train_loader
        for e in range(epochs):

            if args.er_buffer_type=='balanced_batch' and task_id>0 and args.er_size_per_class>0:
                buffer_train=model_container.replay_buffer
                if isinstance(buffer_train, RehearsalMemory):
                    mem_x, mem_y, mem_t = buffer_train.get()
                elif isinstance(buffer_train, BaseTaskSet):
                    mem_x, mem_y, mem_t = buffer_train.get_raw_samples()
                if isinstance(mem_x[0], str):
                    args.tasktype=TaskType.IMAGE_PATH

                train_taskset=_train_loader.dataset
                replay_taskset = TaskSet(x=mem_x, y=mem_y, t=mem_t, target_trsf=train_taskset.target_trsf, trsf=train_taskset.trsf, data_type=args.tasktype)
                replay_dataloader = prepare_dataloader(args, replay_taskset, val_split=0, num_workers=2, shuffle=True)[0]
                def ziped_iterator(loader_1, loader_2):
                    for x_1, y_1, t_1 in loader_1:
                        x_2, y_2, t_2 = next(iter(loader_2))
                        x_2=x_2.squeeze()
                        x=torch.cat([x_1, x_2])
                        y=torch.cat([y_1, y_2])
                        t=torch.cat([t_1, t_2])
                        idx = torch.randperm(len(y))
                        x,y,t=x[idx],y[idx],t[idx]
                        yield x,y,t
                if len(replay_dataloader)<len(_train_loader):
                    replay_dataloader=cycle(replay_dataloader)
                train_loader = ziped_iterator(_train_loader, replay_dataloader)

            if model.optimizer is not None:
                for pg_i, param_group in enumerate(model.optimizer.param_groups):
                    log_wandb({f'train_during/lr_t{task_id}_{pg_i}':param_group["lr"]})
            acc = 0
            # if args.ddp:
            #     train_loader.sampler.set_epoch(e)

            pbar = tqdm(train_loader)
            if args.estimate_compute_regime:
                model.train()
                result = estimate_compute_regime(train_loader, model, epochs, estimate_time=args.estimate_time)
                return model, best_valid_acc, result

            for b_i, (x,y,_) in enumerate(pbar):
                x,y = preprocess_tensor(model, x, y)
                x,y = x.to(args.device), y.to(args.device)
                model.train()
                if len(x.shape)<2:
                    #in the Flower dataset getting x of wrong dimention for some reason (shape (1024))
                    x=x.unsqueeze(0)

                if args.record_flops:
                    if flops_per_batch_fw[len(y)] is None: #keep track of compute per batch size used
                        flops = FlopCountAnalysis(model, x)
                        flops.tracer_warnings('none')
                        flops_per_batch_fw[len(y)] = flops.total()
                        flops_per_batch_bw[len(y)] = flops_per_batch_fw[len(y)]*2
                    flops_meters.add({'flops_task_forward':flops_per_batch_fw[len(y)],'flops_task_backward':flops_per_batch_bw[len(y)], 'flops_task':flops_per_batch_fw[len(y)]+flops_per_batch_bw[len(y)]})
                    if not early_stopped:
                        flops_meters.add({'flops_task_forward_es':flops_per_batch_fw[len(y)],'flops_task_backward_es':flops_per_batch_bw[len(y)], 'flops_task_es':flops_per_batch_fw[len(y)]+flops_per_batch_bw[len(y)]})

                if model_container.args.multihead:
                    raise NotImplementedError
                    y-=exp_state.n_unique_classes[-1] #to make it multi head compatible
                model.zero_grad(set_to_none=True)
                # try:
                start=time.time()
                logits = model(x, task_id, y=y, epoch=e)
                if args.er_buffer_type=='balanced_batch':
                    count=np.bincount(y.cpu(), minlength=logits.shape[1])+1e-10
                    weights =torch.tensor((max(count)/count)/sum((max(count)/count)), device=args.device).float()
                    loss = F.cross_entropy(logits,y, weight=weights)
                else:
                    loss = loss_function(logits,y)
                if not model.optimizer is None and loss.requires_grad:
                    loss.backward()
                    model.optimizer.step()

                end=time.time()
                if time_per_batch is None:
                    time_per_batch=end-start

                acc_c = torch.sum(logits.max(1)[1] == y).float()/len(y)
                acc += acc_c
                pbar.set_description("Acc %s" % (acc.cpu().item()/(b_i+1)))
                if args.debug and args.regime=='sample_ER' and b_i ==5:
                    break
            if lr_anneal and lr_scheduler is not None:
                lr_scheduler.step()
            pbar.close()
            print('model',model_name,'task',task_id, 'train acc: ',acc/len_train_loader, 'epoch: ',e,'\n')
            log_wandb({f'train_during/train_acc_during_t{task_id}':acc/len_train_loader})

            if e%args.test_every==0 and test_loaders is not None:
                test_acc_current_task = test(model, task_id, test_loaders[-1], rank=args.device, args=args)
                print('model',model_name,'task',task_id, 'test acc: ',test_acc_current_task, 'epoch: ',e,'\n')
                log_wandb({f'test_during/test_acc_during_t{task_id}':test_acc_current_task})
            if e%args.validate_every==0:
                valid_acc = test(model, task_id, val_loader, rank=args.device, args=args)
                print('model',model_name,'task',task_id, 'valid acc: ',valid_acc, 'epoch: ',e,'\n')
                log_wandb({f'valid_during/valid_acc_during_t{task_id}':valid_acc})
                if valid_acc>best_valid_acc:
                    n_epochs_without_improvement=0
                    best_valid_acc=valid_acc
                    best_model = model.state_dict()
                    best_model_epoch=e
                else:
                    n_epochs_without_improvement+=1
                    if n_epochs_without_improvement>=args.early_stopping_patience and not early_stopped:
                        early_stopped=1
                        if args.early_stopping:
                            break
                        convergence_epoch=e

        log_wandb({f'convergence_epoch_{task_id}':convergence_epoch})
        log_wandb({f'best_model_epoch_{task_id}':best_model_epoch})
        if epochs>1 and best_model is not None:
            # if best_model is None:
            #     best_model = model.state_dict()
            model.load_state_dict(best_model)
        time_per_epoch=0
        time_per_task=0
        if time_per_batch is not None:
            time_per_epoch=time_per_batch*len_train_loader
            time_per_task=time_per_epoch*epochs
        else:
            time_per_batch=0

        result_compute = UsedFlops(total_flops_bw_epoch=flops_per_epoch_bw,
                                  total_flops_fw_epoch=flops_per_epoch_fw,
                                  total_gigaflops=flops_meters.meters['flops_task'].value()[0]/10e9,
                                  total_gigaflops_bw=flops_meters.meters['flops_task_backward'].value()[0]/10e9,
                                  total_gigaflops_fw=flops_meters.meters['flops_task_forward'].value()[0]/10e9,
                                  total_gigaflops_es=flops_meters.meters['flops_task_es'].value()[0]/10e9,
                                  total_gigaflops_bw_es=flops_meters.meters['flops_task_backward_es'].value()[0]/10e9,
                                  total_gigaflops_fw_es=flops_meters.meters['flops_task_forward_es'].value()[0]/10e9,
                                  walltime_epoch=time_per_epoch,
                                  walltime_task=time_per_task)
        return model, best_valid_acc, result_compute

def learn_task(model_container:ModelContainer_ER, args:ArgsGenerator, model_name:str, task_id:int, train_loader:DataLoader, val_loader:DataLoader,  train_loaders:List[DataLoader]=None, test_loaders:List[DataLoader]=None, params_table:wandb.Table=None, world_size=None):
    # model=model_container.model
    best_valid_acc=0
    best_model = None
    best_result_flops = None
    best_args_model, best_args_classifier = None, None
    #iterate over agruments to validate for per task cross validation
    for cv_run_idx, (args_model, args_classifier, model) in enumerate(model_container.get_model_for_training(task_id)):
        print(f'testing with arguments {args_model} and classifier arguments {args_classifier}')
        if args_model is not None:
            lr_anneal=args_model.lr_anneal
        else:
            lr_anneal = model_container.args.lr_anneal
        #task evel cross validation
        if args.schedule == 'classifier+normal':
            #first trin the classifier and then the feature encoder
            assert args.regime != 'latent_ER'
            #fix feature extractor and finetuen only classifier
            model.freeze_feature_extractor()
            model, valid_acc,result_flops = train_model(model_container, args, model_name, task_id, model, train_loader, val_loader, train_loaders, test_loaders, lr_anneal=lr_anneal)
            #finetune both
            model.freeze_feature_extractor(False)
            model, valid_acc,result_flops = train_model(model_container, args, model_name, task_id, model, train_loader, val_loader, train_loaders, test_loaders, lr_anneal=lr_anneal)
        else:
            if args.n_epochs_task_level_cv is not None and args.task_level_cv and not (args.keep_best_params_after_first_task and task_id>0) :
                epochs=args.n_epochs_task_level_cv
            else:
                epochs=None
            if args.finetuning_only_norms:
                model.freeze_feature_extractor()
                model.unfreeze_bn()
                if model_container.args.encoder_name=='ViT-B/16':
                    if args.unfreeze_input_layer:
                        model.unfreeze_first()
            if model_container.args.encoder_name=='ViT-B/16' and args.freeze_vit_untill_layer is not None:
                model.freeze_vit_untill_layer(args.freeze_vit_untill_layer)
            model, valid_acc, result_flops = train_model(model_container, args, model_name, task_id, model, train_loader, val_loader, train_loaders, test_loaders, epochs=epochs, lr_anneal=lr_anneal)

        flopws_meter_pertask_cv.add([result_flops.total_gigaflops_fw, result_flops.total_gigaflops_bw, result_flops.total_gigaflops, result_flops.total_gigaflops_fw_es,
                                     result_flops.total_gigaflops_bw_es, result_flops.total_gigaflops_es, result_flops.walltime_task])#[:7])
        flopws_meter_pertask_cv.log_flops(task_id)

        if valid_acc>=best_valid_acc:
            best_result_flops=result_flops
            best_model = model.create_checkpoint()
            best_args_model, best_args_classifier = args_model, args_classifier
            best_valid_acc=valid_acc

    if model_container.args.classifier_type in ['slda', 'nmc']:
        return model

    if best_args_model is None:
        best_result_flops=result_flops
        best_model = model.create_checkpoint()
        best_args_model, best_args_classifier = args_model, args_classifier
        best_valid_acc=valid_acc

    model_container.reset_args(best_args_model, best_args_classifier)

    if args.n_epochs_task_level_cv is not None and args.task_level_cv and not (args.keep_best_params_after_first_task and task_id>0):
        #train for total epochs with best hps
        model, _ = model_container.create_model(best_args_model, best_args_classifier)
        model.load_state_dict(model_container.model.state_dict())
        if args.regime=="latent_ER" and model_container.args.classifier_type=='clip_0_shot':
            model.feature_extractor=nn.Identity()
        model.set_optimizer()
        model, best_valid_acc, result_flops = train_model(model_container, args, model_name, task_id, model, train_loader, val_loader, train_loaders, test_loaders)
        best_result_flops=result_flops
        best_model = model.create_checkpoint()
        flopws_meter_pertask_cv.add([best_result_flops.total_gigaflops_fw, best_result_flops.total_gigaflops_bw, best_result_flops.total_gigaflops, best_result_flops.total_gigaflops_fw_es,
                                     best_result_flops.total_gigaflops_bw_es, best_result_flops.total_gigaflops_es, best_result_flops.walltime_task])
        flopws_meter_pertask_cv.log_flops(task_id)

    # if best_args_model is not None:
    print(f"selected arguments model task {task_id}", best_args_model)
    # if best_args_classifier is not None:
    print(f"selected arguments classifier task {task_id}", best_args_classifier)
    # if params_table is not None:
    params_table.add_data(str(task_id), str(best_args_model), str(best_args_classifier))

    #save best parameters to hp database
    if args.use_hp_database and args.task_level_cv:
        if task_id==0 and args.keep_best_params_after_first_task:
            hp_db.log_into_hp_ds(args_global,best_args_model,best_args_classifier)

    if best_result_flops is not None:
        ###log flops (wandb does the accumulation)###
        for n,v in best_result_flops.__dict__.items():
            log_wandb({f'compute/{n}':v})
        #############################################
        flopws_meter_optimal_params.add([best_result_flops.total_gigaflops_fw, best_result_flops.total_gigaflops_bw, best_result_flops.total_gigaflops, best_result_flops.total_gigaflops_fw_es,
                                        best_result_flops.total_gigaflops_bw_es, best_result_flops.total_gigaflops_es, best_result_flops.walltime_task])
        flopws_meter_optimal_params.log_flops(task_id)
        #log per epoch walltime and per epoch compute
        log_wandb({f'compute/walltime_task_t{task_id}':best_result_flops.walltime_task})
        log_wandb({f'compute/walltime_per_epoch_t{task_id}':best_result_flops.walltime_epoch})
        log_wandb({f'compute/flops_per_epoch_t{task_id}': best_result_flops.total_flops_fw_epoch+best_result_flops.total_flops_bw_epoch})

    model_container.init_model()
    if best_model is None:
        best_model = model.create_checkpoint()
    model_container.model.train()
    model_container.model.load_checkpoint(best_model)
    model_container.model.set_optimizer()

    if model_container.args.classifier_type in ['slda', 'nmc']:
        #mnake sure slda and nmc properly store their state dicts
        accuracy=test(model_container, task_id, val_loader, args=args, rank=args.device)
        if not best_valid_acc==accuracy:
            #temporal strange fix for the problem
            model_container.model.load_checkpoint(best_model)
            accuracy=test(model_container, task_id, val_loader, args=args, rank=args.device)
        assert best_valid_acc==accuracy
        # problem with SLDA after second task or so:
        # best_model = model.create_checkpoint()
        # model_container.model.load_checkpoint(best_model)
        # assert test(model, task_id, val_loader, args=args, rank=args.device)==test(model_container, task_id, val_loader, args=args, rank=args.device)
    return model_container.model

def prepare_model_and_scenario(args:ArgsGenerator,args_model:ModelContainer_ER.Options,args_classifier)->Tuple[ModelContainer,_BaseScenario,_BaseScenario]:
    '''
    Prepares the model and the scenarios to train on.
    '''
    if args.regime=='latent_ER':
        scenario, scenario_test = prepare_scenarios(args, args_model)
        n_classes=[scenario[i].nb_classes for i in range(len(scenario))]
        model_container=ModelContainer_ER(args_model, args_global=args, args_classifier=args_classifier, n_classes=n_classes, device=args.device)
        model_container.transforms=None
        model_container.transforms_val=None
        if not args.encode_with_continuum:
            args_model.in_size = int(scenario.dataset[0].shape[-1])
        model_container.init_model()
    elif args.regime=='sample_ER':
        args_model.in_size = int(args.dataset.dataset_info.size[-1])
        model_container=ModelContainer_ER(args_model, args_global=args, args_classifier=args_classifier, n_classes=args.n_classes, device=args.device)
        if args_model.flatten_image:
            args_model.in_size = int(np.prod(args.dataset.dataset_info.size[-1]))
        model_container.init_model()
        scenario, scenario_test = prepare_scenarios(args, args_model, transformations=model_container.transforms, transforms_val=model_container.transforms_val)
        if args.permute_task_order:
            n_classes=[scenario[i].nb_classes for i in range(len(scenario))]
            model_container=ModelContainer_ER(args_model, args_global=args, args_classifier=args_classifier, n_classes=n_classes, device=args.device)
            model_container.init_model()
        if args.finetuning_only_norms:
            model_container.model.freeze_feature_extractor()
            model_container.model.unfreeze_bn()
            if model_container.args.encoder_name=='ViT-B/16':
                if args.unfreeze_input_layer:
                    model_container.model.unfreeze_first()
        if model_container.args.encoder_name=='ViT-B/16':
            if args.freeze_vit_untill_layer is not None:
                model_container.model.freeze_vit_untill_layer(args.freeze_vit_untill_layer)

    return model_container, scenario, scenario_test

def create_per_class_prototypes(x,y,t, n_prototypes_per_class, merge_type='mean', n_samples_per_prototype=None):
    """
        Calculate prototypes for prototype based replay
    """
    n=n_prototypes_per_class
    if merge_type == None:
        return x,y,t
    def split_for_prototype(idxs, n, n_samples_per_prototype):
        random_state = np.random.RandomState(1)
        random_state.shuffle(idxs)
        if n_samples_per_prototype is None:
            #create possibly equal group sizes
            groups_size = [len(idxs) // n + (1 if x < len(idxs) % n else 0)  for x in range (n)]
            #select indicies for group creation without replacement
            groups = [ idxs[sum(groups_size[:max(0,c)]):sum(groups_size[:c])+groups_size[c]] for c in range(len(groups_size)) ]
        else:
            #randomly select indicies fro groups with replacement
            groups = [random_state.choice(idxs,n_samples_per_prototype, replace=True)]
        return groups
    x_, y_, t_ = [],[],[]
    groups_idxs_per_class = [split_for_prototype(np.where(y==c)[0],n,n_samples_per_prototype) for c in np.unique(y)]
    for c_idxs in  groups_idxs_per_class:
        prototypes, label, task = [np.mean(x[i],axis=0) for i in c_idxs], np.concatenate([y[i] for i in c_idxs]), np.concatenate([t[i] for i in c_idxs])
        assert len(np.unique(label))==1
        assert len(np.unique(task))==1
        x_.append(prototypes)
        y_.append([np.unique(label)[0]]*len(prototypes))
        t_.append([np.unique(task)[0]]*len(prototypes))

    return np.concatenate(x_), np.concatenate(y_), np.concatenate(t_)

def main(args:ArgsGenerator, args_model:ModelContainer_ER.Options, args_classifier, **kwargs):
    exp_state=ExperimentState()
    test_loaders_sofar=[]
    train_loaders_sofar=[]
    valid_loaders_sofar=[]
    test_accuracies_past=[]
    valid_accuracies_past=[]
    print(args.device)
    model_container, scenario, scenario_test = prepare_model_and_scenario(args,args_model,args_classifier)
    print(model_container.model)
    logger=Logger(args,model_container,scenario.nb_tasks)
    #wandb table for logging best selected hyperparameters per task
    best_cvparams_table = logger.best_cvparams_table

    #######
    # remapping is necessary for permuted order scenarios
    class_mapping = get_scenario_remapping(scenario)
    model_container.set_mapping(class_mapping)
    for task_id, (train_taskset, test_taskset) in enumerate(zip(scenario, scenario_test)):
        #######
        if args.dataset_name=='MNIST_bckgrndwap':
            #currently the train_taskset only contains transfroms of the current task, the other ones are None
            #current task_set should contain transfroms for old tasks as well when all tasks are being replayed (maybe this should be fixed on the continuum side)
            train_taskset.trsf=[Compose([*scenario.inc_trsf[t],ToTensor()]) for t in range(len(scenario))]
        ########
        exp_state.current_task=task_id
        log_wandb({'task': task_id})

        if isinstance(train_taskset, H5TaskSet):
            #using H5TaskSet was taking to long, probably due to reading from hd5 file,
            # also maybe the issue is the netowrk connection to scratch on mila cluster, maybe putting in into $TEMP(local to the node) should help
            train_taskset = convert_to_task_set(args,train_taskset)
            train_taskset, val_taskset, train_idxs = split_train_val(train_taskset, val_split=args.valid_fraction, valid_k=args.valid_k)
            if model_container.transforms_val is not None:
                if isinstance(model_container.transforms_val, List):
                    transforms_val = Compose(model_container.transforms_val)
                else:
                    transforms_val=model_container.transforms_val
                val_taskset.trsf = transforms_val

        else:
            train_taskset, val_taskset, train_idxs = split_train_val(train_taskset, val_split=args.valid_fraction, valid_k=args.valid_k)
            #set validation transfroms instead of training transfroms for val_taskset
            if model_container.transforms_val is not None:
                if isinstance(model_container.transforms_val, List):
                    if not isinstance(model_container.transforms_val[0], List):
                        transforms_val = Compose(model_container.transforms_val)
                    else:
                        transforms_val = Compose(model_container.transforms_val[task_id])
                else:
                    transforms_val=model_container.transforms_val
                val_taskset.trsf = transforms_val
        #######################################
        # Task similarity logging
        # if args.log_task_similarity and args.regime=='latent_ER':
        #     logger.log_similarity(train_taskset, task_id)
        #######################################
        #for per task hp cv we might want to add some samples from the replay buffer to the validation set
        if args.fraction_buffer_samples_valid>0 and args.er_size_per_class>0 and task_id>0:
            buffer_train, buffer_valid = split_er_buffer(model_container.replay_buffer, args.fraction_buffer_samples_valid, data_type=val_taskset.data_type)#TaskType.IMAGE_PATH if args.regime=='sample_ER' else TaskType.TENSOR)
            #merge buffer_valid into the validation set
            #will take the transfromations fromt the val_taskset
            val_taskset_for_hp_tuning = concat([val_taskset,buffer_valid])
        else:
            buffer_train=model_container.replay_buffer
            val_taskset_for_hp_tuning=val_taskset

        ########################################
        #handle random replay buffer
        if args.er_size_per_class>0 and task_id > 0:
            if args.er_buffer_type=='random':
                train_loader=prepare_randombuffer_train_loader(args, train_taskset=train_taskset, buffer_train=buffer_train)
            elif args.er_buffer_type=='balanced':
                train_loader=prepare_balanced_train_loader(args, buffer_train, current_task_set=train_taskset)
            # elif args.er_buffer_type=='balanced_batch':
            #     train_loader=prepare_balanced_batch_train_loader(args, buffer_train, train_taskset)
            else:
                train_loader,_=prepare_dataloader(args, train_taskset, val_split=0., num_workers=2, shuffle=True)
        else:
            train_loader,_=prepare_dataloader(args, train_taskset, val_split=0., num_workers=2, shuffle=True)
        ########################################

        val_loader,_=prepare_dataloader(args, val_taskset, num_workers=2, val_split=0.)
        val_loader_hp_tuning,_=prepare_dataloader(args, val_taskset_for_hp_tuning, num_workers=2, val_split=0.)
        test_loader,_=prepare_dataloader(args, test_taskset, num_workers=2, val_split=0.)
        test_loaders_sofar.append(test_loader)
        train_loaders_sofar.append(train_loader)
        valid_loaders_sofar.append(val_loader)
        model_container.ready_for_new_task(task_id=task_id, new_classes=scenario[task_id].get_classes(), expand_single_head=args.dataset_name!='MNIST_bckgrndwap')
        if args.debug:
            print(model_container)
        if args.concat_validation_sets:
            val_taskset_for_hp_tuning = concat([loader.dataset for loader in valid_loaders_sofar])
            val_loader_hp_tuning,_=prepare_dataloader(args, val_taskset_for_hp_tuning, num_workers=2, val_split=0.)
        if args.concat_test_sets:
            test_taskset = concat([loader.dataset for loader in test_loaders_sofar])
            test_loader,_=prepare_dataloader(args, val_taskset_for_hp_tuning, num_workers=2, val_split=0.)
        model_container.model = learn_task(model_container, args, model_container.args.encoder_name, task_id, train_loader, val_loader_hp_tuning, train_loaders_sofar, test_loaders_sofar, params_table=best_cvparams_table)
        test_acc =  test(model_container, task_id, test_loader, args=args, rank=args.device)
        valid_acc =  test(model_container, task_id, val_loader, args=args, rank=args.device)
        valid_acc_hp_tuning_set =  test(model_container, task_id, val_loader_hp_tuning, args=args, rank=args.device)
        log_wandb({"test_acc_online":test_acc}, step=('task', task_id))#, prefix='test/')
        log_wandb({"valid_acc_online":valid_acc}, step=('task', task_id))#, prefix='test/')
        log_wandb({"valid_acc_online(hp_tuning)":valid_acc_hp_tuning_set}, step=('task', task_id))#, prefix='test/')
        log_wandb({f"test_acc_best_t{task_id}":test_acc}, step=('task', task_id))#, prefix='test/')
        log_wandb({f"valid_acc_best_t{task_id}":valid_acc}, step=('task', task_id))#, prefix='test/')
        test_accuracies_past.append(test_acc)
        valid_accuracies_past.append(valid_acc)
        #add to replay buffer
        if args.er_size_per_class>0:
            x,y,t = scenario[task_id].get_raw_samples()
            if len(y.shape)>1:
                y=y[0]
                t=t[0]
            #only add samples from train set (train_idxs), not the validation set to the replay buffer
            x,y,t = x[train_idxs],y[train_idxs],t[train_idxs]
            print(len(train_idxs))
            if args.er_with_prototypes:
                #create er_size_per_class prototypes
                if args.size_of_random_prototype_factor is not None:
                    size_of_random_prototype = int(np.mean(np.bincount(scenario[task_id]._y))/args.er_size_per_class)
                    size_of_random_prototype*=args.size_of_random_prototype_factor
                    size_of_random_prototype=int(size_of_random_prototype)
                else:
                    size_of_random_prototype=None
                x,y,t = create_per_class_prototypes(x,y,t, n_prototypes_per_class=args.er_size_per_class, n_samples_per_prototype=size_of_random_prototype)
            model_container.add_to_buffer((x,y,t))
        logger.log_results(args,task_id,model_container,test_loaders_sofar,valid_loaders_sofar,test_accuracies_past,valid_accuracies_past,best_cvparams_table)
        if args.stop_after_first_task and task_id==0:
            break
        if args.reinit_between_tasks:
            model_container.reinit_model()
        if args_global.save_final_model:
            ckpt_name=f'{model_container.args.encoder_name.replace("/","_")}_trained_{args_global.dataset_name}_task{task_id}_{args_global.n_tasks}tasks_er_size_{args_global.er_size_per_class}'
            if args_model.checkpoint is not None:
                ckpt_name+=args_model.checkpoint
            torch.save({
                'model_state_dict': model_container.model.state_dict(),
            }, args_global.weights_path+f'/{ckpt_name}.ckpt')

    flopws_meter_pertask_cv.log_flops()
    flopws_meter_optimal_params.log_flops()
    logger.close()

if __name__== "__main__":
    parser = ArgumentParser()
    parser.add_arguments(ArgsGenerator, dest="Global")
    parser.add_arguments(ModelContainer.Options, dest="Model")
    parser.add_arguments(Classifier_options, dest="Classifier")
    args = parser.parse_args()
    args_global:ArgsGenerator = args.Global
    args_model = args.Model
    if args_model.classifier_type=='fc':
        args_classifier = args.Classifier.CLS_NN
    elif args_model.classifier_type=='logistic_regression':
        args_classifier = args.Classifier.CLS_LogReg
    elif args_model.classifier_type=='BiT_classifier':
        args_classifier = args.Classifier.CLS_BiT
    elif args_model.classifier_type=='random_forrest':
        args_classifier = args.Classifier.CLS_RF
    elif args_model.classifier_type=='clip_0_shot':
        args_classifier = args.Classifier.CLP_0
    elif args_model.classifier_type == 'knn':
        args_classifier = args.Classifier.CLS_KNN
    elif args_model.classifier_type == 'weightnorm':
        args_classifier = args.Classifier.CLS_WeightNorm
    elif args_model.classifier_type == 'slda':
        args_classifier = args.Classifier.CLS_SLDA
    elif args_model.classifier_type == 'nmc':
        args_classifier = args.Classifier.CLS_NMC
        args_global.device = 'cpu'
    else:
        args_classifier=None

    wandb_project = args_global.wandb_project if not args_global.debug else 'test'
    if args_global.group_name=='':
        args_global.group_name = wandb.util.generate_id()
    # assert is_connected()#:
    if not is_connected():
        print('no internet connection. Going in dry')
        # os.environ['WANDB_MODE'] = 'dryrun'
        wandb_mode="offline"
    else:
        wandb_mode=None

    def start_experiment(run):
        wandb.config.update(args_global, allow_val_change=True)
        wandb.config.update(args_model, allow_val_change=True)
        if args_classifier is not None:
            wandb.config.update(args_classifier, allow_val_change=True)
        main(args_global, args_model, args_classifier, world_size=torch.cuda.device_count())
        run.finish()

    key=f'{args_global.md5}_{args_model.md5}_{args_classifier.md5}'
    #databse for hyperparameter saving: hp dataset is job specific
    if args_global.use_hp_database:
        if 'SLURM_TMPDIR' in os.environ:
            args_global.hp_db_path = os.environ.get('SLURM_TMPDIR')
        hp_db = TinyDB_hp_tracker(args_global.hp_db_path+f'/hp_database_{args_global.dataset_name}.json', key=key)
    args_global.group_name = args_global.generate_group_name(args_model.md5+args_classifier.md5)
    set_seed(manualSeed=args_global.seed)
    if args_global.n_task_order_permutations>0:
        args_global.permute_task_order=1
        for run_i in range(args_global.n_task_order_permutations):
            flopws_meter_pertask_cv, flopws_meter_optimal_params = init_flops_meters()
            args_global.generate_task_order(run_i)
            if args_global.use_hp_database:
                success,args_global,args_model,args_classifier=hp_db.load_hps_from_db(args_global,args_model,args_classifier)
            run = wandb.init(project=wandb_project, mode=wandb_mode, notes=args_global.wandb_notes, group=args_global.group_name, settings=wandb.Settings(start_method="fork"), reinit=False, entity=args_global.wandb_entity)
            wandb.save('datasets.py')
            start_experiment(run)
    else:
        flopws_meter_pertask_cv, flopws_meter_optimal_params = init_flops_meters()
        if args_global.use_hp_database:
            success,args_global,args_model,args_classifier=hp_db.load_hps_from_db(args_global,args_model,args_classifier)
        run = wandb.init(project=wandb_project, mode=wandb_mode, notes=args_global.wandb_notes, settings=wandb.Settings(start_method="fork"), reinit=False, entity=args_global.wandb_entity)
        wandb.save('datasets.py')
        start_experiment(run)