import torch
from torch import nn
from random import shuffle
import numpy as np
import matplotlib.pyplot as plt

class ff_network(nn.Module):

    def __init__(self):
        """
        Fully Connected layers
        """
        super(ff_network, self).__init__()

        self.semnet = nn.Sequential(
            nn.Linear(17, 50), # 17 properties
            nn.ReLU(),
            nn.Linear(50, 10), # very small network for tests
            nn.ReLU(),
            nn.Linear(10, 1)
        )

    def forward(self, x):
        """
        Pass throught network
        """
        res = self.semnet(x)

        return res

def train_model(data_train0, data_test0, data_train1, data_test1, lr_enc, batch_size):

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model_semnet = ff_network().to(device)
    optimizer_predictor = torch.optim.Adam(model_semnet.parameters(), lr=lr_enc)

    data_train0=torch.tensor(data_train0, dtype=torch.float).to(device)
    data_test0=torch.tensor(data_test0, dtype=torch.float).to(device)

    data_train1=torch.tensor(data_train1, dtype=torch.float).to(device)
    data_test1=torch.tensor(data_test1, dtype=torch.float).to(device)

    test_loss_total=[]
    for epoch in range(500): # should be much larger, with good early stopping criteria
        small_train_len=min(len(data_train0),len(data_train1))
        for batch_iteration in range(int(small_train_len/batch_size)):
            idx = torch.randint(0, len(data_train0), (batch_size,))
            data_train0_samples = data_train0[idx]

            idx = torch.randint(0, len(data_train1), (batch_size,))
            data_train1_samples = data_train1[idx]

            calc_properties0 = model_semnet(data_train0_samples)
            criterion = torch.nn.MSELoss()
            curr_pred=torch.tensor([0] * batch_size, dtype=torch.float).to(device)
            real_loss0 = criterion(calc_properties0, curr_pred)
            loss0 = torch.clamp(real_loss0, min = 0., max = 50000.).double()

            calc_properties1 = model_semnet(data_train1_samples)

            criterion = torch.nn.MSELoss()
            curr_pred=torch.tensor([1] * batch_size, dtype=torch.float).to(device)
            real_loss1 = criterion(calc_properties1, curr_pred)
            loss1 = torch.clamp(real_loss1, min = 0., max = 50000.).double()

            loss=loss0+loss1
            optimizer_predictor.zero_grad()
            loss.backward()
            optimizer_predictor.step()

        # calculate train set
        model_semnet.eval()

        calc_properties0 = model_semnet(data_train0[1:3000])
        criterion = torch.nn.MSELoss()

        curr_pred=torch.tensor([0] * len(data_train0), dtype=torch.float).to(device)

        real_loss0 = criterion(calc_properties0, curr_pred)

        real_loss_train_num0=real_loss0.detach().cpu().numpy()

        calc_properties1 = model_semnet(data_train1)
        criterion = torch.nn.MSELoss()
        curr_pred=torch.tensor([1] * len(data_train1), dtype=torch.float).to(device)
        real_loss1 = criterion(calc_properties1, curr_pred)
        real_loss_train_num1=real_loss1.detach().cpu().numpy()

        # calculate test set
        calc_properties0 = model_semnet(data_test0[1:3000])
        criterion = torch.nn.MSELoss()
        curr_pred=torch.tensor([0] * len(data_test0), dtype=torch.float).to(device)
        real_loss0 = criterion(calc_properties0, curr_pred)
        real_loss_test_num0=real_loss0.detach().cpu().numpy()

        calc_properties1 = model_semnet(data_test1)
        criterion = torch.nn.MSELoss()
        curr_pred=torch.tensor([1] * len(data_test1), dtype=torch.float).to(device)
        real_loss1 = criterion(calc_properties1, curr_pred)
        real_loss_test_num1=real_loss1.detach().cpu().numpy()

        test_loss_total.append(np.mean(real_loss_test_num0)+np.mean(real_loss_test_num1))

        if epoch%10==0:
            info_str='epoch: '+str(epoch)+' - totalloss: ',np.mean(real_loss_train_num0)+np.mean(real_loss_train_num1) ,'; l1/l2: '+str(np.mean(real_loss_train_num0))+'/'+str(np.mean(real_loss_train_num1))+'; total: ',np.mean(real_loss_test_num0)+np.mean(real_loss_test_num1) ,' ts1/ts2: '+str(np.mean(real_loss_test_num0))+'/'+str(np.mean(real_loss_test_num1))
            print('train_model: ',info_str)
        model_semnet.train()

        if len(test_loss_total)>30: # early stopping
            if (test_loss_total[-10]-test_loss_total[-1])<=0 and (test_loss_total[-11]-test_loss_total[-1])<=0 and (test_loss_total[-12]-test_loss_total[-1])<=0:
                break

    plt.plot(test_loss_total)
    plt.pause(0.05)
    plt.show

    return True

def train_nn_one_instance(data_0,data_1,model_semnet):
    batch_size=100
    lr_enc=0.01

    model_semnet.train()

    train_valid_test_size=[0.5, 0.5, 0.0]
    x = [i for i in range(len(data_0))]  # random shuffle input
    shuffle(x)
    data_0 = data_0[x]
    idx_traintest=int(len(data_0)*train_valid_test_size[0])
    data_train0=data_0[0:idx_traintest]
    data_test0=data_0[idx_traintest:]

    x = [i for i in range(len(data_1))]  # random shuffle input
    shuffle(x)
    data_1 = data_1[x]
    idx_traintest=int(len(data_1)*train_valid_test_size[0])
    data_train1=data_1[0:idx_traintest]
    data_test1=data_1[idx_traintest:]

    print('train_nn_one_instance - start training')
    bestval_no=train_model(data_train0, data_test0, data_train1, data_test1, lr_enc, batch_size)

    print('train_nn_one_instance - Calculate ROC')

    return bestval_no

def calculate_ROC(data_0,data_1,model_semnet):
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    model_semnet.eval()
    data_0=torch.tensor(data_0, dtype=torch.float).to(device)
    data_1=torch.tensor(data_1, dtype=torch.float).to(device)
    future_data_0=model_semnet(data_0).detach().cpu().numpy()
    future_data_1=model_semnet(data_1).detach().cpu().numpy()

    print(future_data_0.shape)
    print(future_data_1.shape)

    all_pred=np.concatenate([future_data_0,future_data_1])
    corr_answers=np.array([0] * len(future_data_0) + [1] * len(future_data_1))
    idx1=np.flip(np.argsort(all_pred,axis=0)) # indices from strongest predicted connection to shortest one

    print(corr_answers.shape)
    corr_answers_sorted=corr_answers[idx1]

    xpos=[0]
    ypos=[0]
    ROC_vals=[]
    for ii in range(len(corr_answers_sorted)):
        if corr_answers_sorted[ii]==1:
            xpos.append(xpos[-1])
            ypos.append(ypos[-1]+1)
        if corr_answers_sorted[ii]==0:
            xpos.append(xpos[-1]+1)
            ypos.append(ypos[-1])
            ROC_vals.append(ypos[-1])

    ROC_vals=np.array(ROC_vals)/max(ypos)
    ypos=np.array(ypos)/max(ypos)
    xpos=np.array(xpos)/max(xpos)

    plt.plot(xpos, ypos)
    plt.show()

    AUC=sum(ROC_vals)/len(ROC_vals)
    print('Area Under Curve: ', AUC,'\n\n\n')

    return future_data_0, future_data_1

def train_nn(all_data_0, all_data_1, prediction_distance, start_year):
    for y in range(len(all_data_0[1:3])):
        print('train_nn - year ',start_year+y+2)

        model_semnet = ff_network()
        train_nn_one_instance(all_data_0[y],all_data_1[y],model_semnet)
        print('train_nn - finished - year ',start_year+y+2)
        print('Calculate ROC & AUC')

        if y+prediction_distance<len(all_data_0):
            future_data_0, future_data_1 = calculate_ROC(all_data_0[y+prediction_distance],all_data_1[y+prediction_distance],model_semnet)

    return True


#审查清理文本字符串
import unicodedata
import sys
s = 'pýtĥöñ\fis\tawesome\r\n'
#创建转换表格
remap={
    ord('\t'): ' ',
    ord('\f'): ' ',
    ord('\r'): None,#这里的ord好像是必须的
}
a=s.translate(remap)
print(a)

cmb_chrs=dict.fromkeys(c for c in range(sys.maxunicode) if unicodedata.combining(chr(c)))
print(cmb_chrs)
b=unicodedata.normalize('NFD',a)
print(b)
print(b.translate(cmb_chrs))

digitmap={c : ord('0')+unicodedata.digit(chr(c)) for c in range(sys.maxunicode) if unicodedata.category(chr(c))=='Nd'}
print(digitmap)
print( '\u0661\u0662\u0663'.translate(digitmap))

b=unicodedata.normalize('NFD',a)
c=b.encode('ascii','ignore').decode('ascii')
print(c)

#使用str.replace处理简单字符是最快的，处理复杂字符则使用translate


# -*- coding: utf-8 -*-
"""
Created on Sat Nov 27 11:56:13 2021

@author: Checkout
"""

import requests
import json
import pandas as pd
if __name__=="__main__":
    url = 'http://localhost:5000/testApi'
    r = requests.post(url,json={"side": 8,
                                "bomb_count": 6,
                                "start": [2,2],
                                "goal":[7,7]})
    print(r.content)
    data = json.loads(r.text)
    pd.DataFrame.from_dict(eval(data))


class Solution:
    def uniquePathsIII(self, grid: List[List[int]]) -> int:
        row = len(grid)
        col = len(grid[0])
        self.tot_count = 0
        self.output = 0
        visited = set()

        def helper(i,j,count):
            if i >= row or i < 0 or j >= col or j < 0:
                return

            if grid[i][j] == -1 or grid[i][j] == 1 or (i,j) in visited:
                return

            if grid[i][j] == 2:
                if count == self.tot_count:
                    self.output +=1
                return

            visited.add((i,j))

            direction = [(-1,0),(1,0),(0,-1),(0,1)]
            for x,y in direction:
                helper(i + x,j + y,count + 1)

            visited.remove((i,j))
            return

        # count total number of empty squares we can walk over
        for i in range(row):
            for j in range(col):
                if grid[i][j] == 0:
                    self.tot_count +=1

        for i in range(row):
            for j in range(col):
                if grid[i][j] == 1:
                    helper(i - 1,j,0)
                    helper(i + 1,j,0)
                    helper(i,j - 1,0)
                    helper(i,j + 1,0)
                    break

        return self.output


"""
 Pygame base template for opening a window

 Sample Python/Pygame Programs
 Simpson College Computer Science
 http://programarcadegames.com/
 http://simpson.edu/computer-science/

 Explanation video: http://youtu.be/vRB_983kUMc
"""

import pygame
import random

# Define some colors
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
GREEN = (0, 255, 0)
RED = (255, 0, 0)
BLUE = (0, 0, 255)
colors = [BLACK, GREEN, RED, BLUE]
def random_color():
	return random.coice(colors)
pygame.init()

# Set the width and height of the screen [width, height]
screen_width = 700
screen_height = 500
screen = pygame.display.set_mode((screen_width, screen_height))

x_speed = 3
y_speed = 3

pygame.display.set_caption("Documentation Hunt")

class Circle():
	def __init__(self, mouse_position, Circle_size):
		self.mouse_position = mouse_position
		x_pos = mouse_position[0]
		y_pos = mouse_position[1]
		self.x_pos = x_pos
		self.y_pos = y_pos
		self.Circle_size = Circle_size
	def draw(self):
		global screen
		global RED
		pygame.draw.circle(screen,RED, self.mouse_position,self.Circle_size)
	def move(self, x_pos, y_pos, y_speed, x_speed):
		self.x_speed = x_speed
		self.y_speed = y_speed
		self.x_pos = x_pos								#make a list of positions?
		self.y_pos = y_pos

		if self.x_pos >= screen_width - self.Circle_size or self.x_pos < self.Circle_size:
			self.x_speed = x_speed * -1

		if self.y_pos >= screen_height - self.Circle_size or self.y_pos < self.Circle_size:
			self.y_speed = self.y_speed * -1

		self.x_pos += self.x_speed
		self.y_pos += self.y_speed

# Loop until the user clicks the close button.
done = False

# Used to manage how fast the screen updates
clock = pygame.time.Clock()

circle_list = []

screen.fill(WHITE)

# -------- Main Program Loop -----------
while not done:
	# --- Main event loop
	for event in pygame.event.get():
		if event.type == pygame.QUIT:
			done = True

	# --- Game logic should go here
	pressed = pygame.mouse.get_pressed()
	if pressed[0] == 1:			  #checks if mouse click is happening
		print("Here!")
		screen.fill(WHITE)
		mouse_position = pygame.mouse.get_pos()			#takes position where mouse is pressed
		circle = Circle(mouse_position,20)					#makes circle object
		circle_list.append(circle)
		for circle in circle_list:
			circle.draw()								#draws circles
			circle.move(1,1,1,1)

	# --- Screen-clearing code goes here

	# Here, we clear the screen to white. Don't put other drawing commands
	# above this, or they will be erased with this command.

	# If you want a background image, replace this clear with blit'ing the
	# background image

	# --- Drawing code should go here

	# --- Go ahead and update the screen with what we've drawn.
	pygame.display.flip()

	# --- Limit to 60 frames per second
	clock.tick(60)

# Close the window and quit.
pygame.quit()
exit() # Needed when using IDLE


import numpy as np
import pandas as pd
from tqdm import tqdm
import itertools, argparse, os, sys, random, logging, config, torch, torchvision, utils
import torchvision.models as models
import torch.nn as nn

def run_inference(args, model, input_data, device):

	model.eval()
	softmax = nn.Softmax(dim=1)

	starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)

	inference_time_list = []

	with torch.no_grad():
		for i in tqdm(range(args.n_rounds)):

			starter.record()

			_ = model(input_data)

			ender.record()
			torch.cuda.synchronize()
			curr_time = starter.elapsed_time(ender)

			print(curr_time)

			inference_time_list.append(curr_time)

			torch.cuda.empty_cache()

	results = {"inference_time": inference_time_list}

	return results

def save_result(result, save_path):
	df = pd.DataFrame(np.array(list(result.values())).T, columns=list(result.keys()))
	df.to_csv(save_path, mode='a', header=not os.path.exists(save_path) )

def extracting_inference_time_backbone(args, model, input_data, input_dim, dim, save_path, device):

	result = run_inference(args, model, input_data, device)

	save_result(result, save_path)

def main(args):

	model_save_path = os.path.join(config.DIR_NAME, "models", config.dataset_name, config.model_name,
		"%s_backbone_model_%s_%s.pth"%(args.distortion_type, config.model_name, args.model_id))

	inference_time_path = os.path.join(config.DIR_NAME, "inference_data", args.dataset_name, args.model_name,
		"inference_time_backbone_id_%s_final_final.csv"%(args.model_id))

	device = torch.device('cuda' if (torch.cuda.is_available() and args.cuda) else 'cpu')

	n_classes = 257

	#Load the trained backbone DNN model.
	model = models.mobilenet_v2(pretrained=True)
	model.classifier[1] = nn.Linear(1280, n_classes)

	model.load_state_dict(torch.load(model_save_path, map_location=device)["model_state_dict"])
	model = model.to(device)

	dummy_input = torch.randn(1, 3, 224,224, dtype=torch.float).to(device)

	for _ in range(10):
		_ = model(dummy_input)

	extracting_inference_time_backbone(args, model, dummy_input, args.input_dim, args.dim, inference_time_path, device)

if (__name__ == "__main__"):

	parser = argparse.ArgumentParser(description='UCB using MobileNet')
	parser.add_argument('--model_id', type=int, default=config.model_id, help='Model Id.')
	parser.add_argument('--dataset_name', type=str, default=config.dataset_name, help='Dataset Name.')
	parser.add_argument('--model_name', type=str, default=config.model_name, help='Model name.')
	parser.add_argument('--cuda', type=bool, default=config.cuda, help='Cuda ?')
	parser.add_argument('--seed', type=int, default=config.seed, help='Seed.')
	parser.add_argument('--batch_size_train', type=int, default=config.batch_size_train, help='Size of train batch.')
	parser.add_argument('--split_ratio', type=float, default=config.split_ratio, help='Split Ratio')
	parser.add_argument('--distortion_prob', type=float, default=1)
	parser.add_argument('--distortion_type', type=str, default=config.distortion_type, help='Distortion Type.')
	parser.add_argument('--input_dim', type=int, default=config.input_dim, help='Input Dim. Default: %s'%config.input_dim)
	parser.add_argument('--dim', type=int, default=config.dim, help='Dim. Default: %s'%(config.dim))
	parser.add_argument('--n_rounds', type=float, default=10000)

	args = parser.parse_args()
	main(args)


from flask import Flask, render_template, request, redirect, url_for

app = Flask(__name__)

@app.route("/")
def home():
  return render_template("index.html")

@app.route("/result", methods=['POST'])
def result():
    num1 = request.form['num1']
    num2 =request.form['num2']

    num1 = float(num1)
    num2 = float(num2)
    result = 0
    operation = request.form['operation']
    if operation == 'add':
      result = num1 + num2
      content = 'sum'
    elif operation == 'subtract':
      result = num1 - num2
      content = 'difference'
    elif operation == 'multiply':
      result = num1 * num2
      content = 'product'
    elif operation == 'divide':
      content = 'quotient'
      if num2 != 0:
        result = num1 / num2
      else:
        result = "Error: Division by zero"
    return render_template('result.html', result=result, content=content)

if __name__ == "__main__":
  app.run(debug=True)


from os.path import dirname, join as os_path_join, exists, isfile
from os import listdir, makedirs, remove

import csv
from datetime import date

from xlsxwriter.workbook import Workbook

class Logger:
    '''Logger which create an ".xlsx" file which could be set with:

    Parameters:
        logging_name in order to create a file with the following format:
            TodayDate_logging_name_xxxx.xlsx (xxxx is interger)

        logging_folder relative path from logger.py folder path
        where create the log files
        (default: create "logging" folder in project root)

        columns_width in order to manage the width of each column
        (default: not column width modification)
            [40, 10]  ->  set the first colum to 40, the second to 10...

    Functions:
        add_lines_to_logging_file(list):
            each list's items are respectively write in columns of a new row

        stop_logging(sheet_name):
            when all needed data are stored, create the xlsx file
            with a sheet named 'sheet_name'
            (default: default excel file sheet name)

        set_logging_name(logging_name):
            in order to change the logging file name
            after Logger instantiation

        set_column_width(columns_width):
            in order to change excel sheet columns width
            after Logger instantiation
    '''

    def __init__(self,
                 logging_name='multimeter_Measurements',
                 logging_folder="../../logging",
                 columns_width=[]):
        self._logging_folder = logging_folder
        self._logging_logging_name = '_' + logging_name + '_'
        self._set_logging_folder()
        self._generate_logging_file_path()
        self.set_column_width(columns_width)

    def add_lines_to_logging_file(self, data: list):
        with open(
                self._logging_file_path,
                'a',
                newline='',
                encoding='utf8') as file:
            writer = csv.writer(file)
            writer.writerow(data)

    def stop_logging(self, sheet_name=None):
        '''After logging, convert csv to xlsx'''

        csvfile = self._logging_file_path
        # Create Workbook with worksheet named 'sheet_name'
        self._workbook = Workbook(csvfile[:-4] + '.xlsx')
        worksheet = self._workbook.add_worksheet(name=sheet_name)
        # Set cells format
        cell_format = self._set_cell_format()
        cell_format_ok = self._set_cell_format('green')
        cell_format_nok = self._set_cell_format('red')
        cell_format_board_number = self._set_cell_format('cyan')
        # Update column width
        for colum_index, width in enumerate(self._columns_width):
            worksheet.set_column(colum_index, colum_index, width)
        # Check if csv file has been created
        if not exists(csvfile):
            return
        # Write csv in xlsx
        with open(csvfile, 'rt', encoding='utf8') as file:
            reader = csv.reader(file)
            for r, row in enumerate(reader):
                for c, col in enumerate(row):
                    if col == 'Test OK':
                        choice_format = cell_format_ok
                    elif col == 'Test NOK' or col == 'Exception occurs':
                        choice_format = cell_format_nok
                    elif col == 'Board number':
                        choice_format = cell_format_board_number
                    else:
                        choice_format = cell_format
                    worksheet.write(r, c, col, choice_format)
        self._workbook.close()
        # Delete temporary csv file
        remove(csvfile)
        # Update logging file's name
        self._generate_logging_file_path()

    def set_logging_name(self, logging_name: str):
        self._logging_logging_name = '_' + logging_name + '_'
        self._generate_logging_file_path()

    def set_column_width(self, columns_width: list = []):
        self._columns_width = columns_width

    def _set_cell_format(self, bg_color: str = ''):
        cell_format = self._workbook.add_format()
        cell_format.set_align('right')
        cell_format.set_align('vcenter')
        cell_format.set_indent(1)
        if not bg_color == '':
            cell_format.set_bg_color(bg_color)
        return cell_format

    def _files_in_folder(self, folder_path: str) -> list:
        '''Return the list of all files in folder_path and its subfolder'''
        files_in_folder_path = []
        # Loop on all files or directories inside folder_path
        for file_or_dir in listdir(folder_path):
            file_or_dir_path = os_path_join(folder_path, file_or_dir)
            # If file_or_dir is a file
            if isfile(file_or_dir_path):
                # Add it to the files list
                files_in_folder_path.append(file_or_dir)
            else:
                # Loop inside the subfolder to find all files
                files_in_subfolder = self._files_in_folder(file_or_dir_path)
                # Add relative path from folder_path
                for index, file in enumerate(files_in_subfolder.copy()):
                    files_in_subfolder[index] = file_or_dir + "/" + file
                # Add this files to the files list
                files_in_folder_path += files_in_subfolder
        return files_in_folder_path

    def _set_logging_folder(self):
        module_dir = dirname(__file__)
        logging_folder_path = os_path_join(
            module_dir,
            self._logging_folder)
        self._check_folder_exist_or_create(logging_folder_path)
        self._logging_folder_path = logging_folder_path

    def _get_logging_files_list(self):
        logging_files_list = self._files_in_folder(
            self._logging_folder_path)
        return logging_files_list

    def _generate_logging_file_path(self) -> None:
        '''Generate a new default logging file path'''
        logging_files_list = self._get_logging_files_list()
        today = date.today().strftime('%Y-%m-%d')
        for i in range(10000):
            file_name = (
                today
                + self._logging_logging_name
                + self._file_number_padding(i)
                )
            file_name_csv = file_name + ".csv"
            file_name_xlsx = file_name + ".xlsx"
            if (
                    file_name_csv not in logging_files_list
                    and
                    file_name_xlsx not in logging_files_list):
                self._logging_file_path = os_path_join(
                    self._logging_folder_path,
                    file_name_csv)
                break

    def _file_number_padding(self, index: int) -> str:
        if index < 10:
            return "000" + str(index)
        elif index < 100:
            return "00" + str(index)
        elif index < 1000:
            return "0" + str(index)
        else:
            return str(index)

    def _check_folder_exist_or_create(self, dir: str):
        '''
        From an absolute path, create the folder if it doesn't already exist
        '''
        # Check if this folder exist
        if not exists(dir):
            # Created the folders
            makedirs(dir)


import argparse
import boto3
import os
import subprocess
import sys
import tempfile

def zip(src_folder, zip_name):
    """Zip all the files and folders in src_folder.

    Note that the zip command is run _within_ src_folder so the folder,
    itself, will not appear in the zip.

    Args:
        src_folder (string): Folder to zip.
        zip_name (string): Name of output zip file.

    Returns:
        (bool): True on success.
    """

    #args = ('/usr/bin/zip', '--symlinks', '-r', '-q', zip_name, '.')
    args = ('/usr/bin/zip', '-r', '-q', zip_name, '.')
    popen = subprocess.Popen(args, cwd=src_folder, stdout=subprocess.PIPE)
    exit_code = popen.wait()
    output = popen.stdout.read()
    if not exit_code == 0:
        print(str(output))
        return False
    return True

def upload_to_s3(session, zip_file, bucket):
    """Upload the zip file to the given S3 bucket.

    Args:
        session (Session): Boto3 Session.
        zip_file (string): Name of zip file.  The name (after any path is stripped) is used as the key.
        bucket (string): Name of bucket to use.
    """
    key = os.path.basename(zip_file)
    s3 = session.client('s3')
    try:
        s3.create_bucket(Bucket=bucket)
    except s3.exceptions.BucketAlreadyOwnedByYou:
        pass # Only us-east-1 will not throw an exception if the bucket already exists
    s3.put_object(Bucket=bucket, Key=key, Body=open(zip_file, 'rb'))

def setup_parser():
    parser = argparse.ArgumentParser(
        description='Script for deploying lambda functions to S3.')
    parser.add_argument(
        '--aws-credentials', '-a',
        metavar = '<file>',
        default = os.environ.get('AWS_CREDENTIALS'),
        type = argparse.FileType('r'),
        help = 'File with credentials for connecting to AWS (default: AWS_CREDENTIALS)')
    parser.add_argument(
        '--upload-only', '-u',
        action = 'store_const',
        const = True,
        default = False,
        help = 'Don\'t re-create the zip.  Just upload to S3.')
    parser.add_argument(
        'source_folder',
        help = 'Folder containing lambda handlers and the virtualenv.')
    parser.add_argument(
        'zip_file',
        help = 'source_folder will be zipped and stored in this filename.')
    parser.add_argument(
        'bucket',
        help = 'Name of S3 bucket to upload lambdas to.')

    return parser

if __name__ == '__main__':
    parser = setup_parser()
    args = parser.parse_args()

    cwd = os.getcwd()

    # Build an absolute path to the zip file if one not provided.
    zip_file_abs_path = os.path.join(cwd, args.zip_file)

    if not args.upload_only:
        if not zip(args.source_folder, zip_file_abs_path):
            sys.exit(1)

    if args.aws_credentials is None:
        # This allows aws roles to be used to create sessions.
        session = boto3.session.Session()
    else:
        # moved so the script can be independent if not passing credentials
        from lambdautils import create_session
        session = create_session(args.aws_credentials)
    upload_to_s3(session, args.zip_file, args.bucket)

    print('Done.\n')


import h5py
import json
import pandas as pd

import cellranger.analysis.singlegenome as cr_sg_analysis
import cellranger.analysis.multigenome as cr_mg_analysis
import cellranger.webshim.constants.gex as gex_constants

class SampleDataPaths:
    def __init__(self, summary_path=None, barcode_summary_path=None, analysis_path=None,
                 vdj_clonotype_summary_path=None,
                 vdj_barcode_support_path=None):
        self.summary_path = summary_path
        self.barcode_summary_path = barcode_summary_path
        self.analysis_path = analysis_path
        self.vdj_clonotype_summary_path = vdj_clonotype_summary_path
        self.vdj_barcode_support_path = vdj_barcode_support_path

class SampleData:
    def __init__(self, sample_properties, sample_data_paths, plot_preprocess_func):
        self.summary = load_metrics_summary(sample_data_paths.summary_path)

        # Not guaranteed to exist for all pipelines
        self.num_cells = self.summary.get('filtered_bcs_transcriptome_union')

        self.analysis, self.original_cluster_sizes = load_analysis(sample_data_paths.analysis_path,
                                                                   plot_preprocess_func,
                                                                   sample_properties)

        self.barcode_summary = load_barcode_summary(sample_data_paths.barcode_summary_path)

        self.vdj_clonotype_summary = pd.read_csv(sample_data_paths.vdj_clonotype_summary_path) if sample_data_paths.vdj_clonotype_summary_path else None

        self.vdj_barcode_support = pd.read_csv(sample_data_paths.vdj_barcode_support_path, usecols=[1]) if sample_data_paths.vdj_barcode_support_path else None

def load_metrics_summary(summary_path):
    if summary_path is None:
        return None
    with open(summary_path) as f:
        return json.load(f)

def load_analysis(base_dir, plot_preprocess_func, sample_properties):
    """ Returns (analysis_object, original_cluster_sizes) """
    if base_dir is None:
        return None, None

    if len(sample_properties['genomes']) == 1:
        analysis = cr_sg_analysis.SingleGenomeAnalysis.load_default_format(base_dir)
    else:
        analysis = cr_mg_analysis.MultiGenomeAnalysis.load_default_format(base_dir)

    if analysis is None:
        return None, None

    # Subsample barcodes for plotting purposes
    original_cluster_sizes = None
    if analysis is not None and isinstance(analysis, cr_sg_analysis.SingleGenomeAnalysis):
        original_cluster_sizes = analysis.get_cluster_sizes()
        if analysis.matrix.bcs_dim > gex_constants.MAX_WEBSHIM_BCS_DIM:
            analysis.subsample_bcs(gex_constants.MAX_WEBSHIM_BCS_DIM)

    return plot_preprocess_func(analysis, sample_properties), original_cluster_sizes

def load_barcode_summary(h5_path):
    if h5_path is None:
        return None
    return h5py.File(h5_path, 'r')


import tensorflow as tf

def trainmodel(model,xdata,ydata,args):
    NUMBER_EPOCHS = args['nepoch']
    filepath=args['ckpt']
    BATCH_SIZE=args['batchsize']
    validation_split=args['valsplit']

    model_checkpoint_callback=tf.keras.callbacks.ModelCheckpoint(
        filepath,
        monitor="val_loss",
        verbose=1,
        save_best_only=True,
        save_weights_only=False,
        mode="min",
        save_freq="epoch",
        options=None
    )
    hist = model.fit(x=xdata,
                     y=ydata,
                     epochs=NUMBER_EPOCHS,
                     batch_size=BATCH_SIZE,
                     validation_split=validation_split,
                     verbose=1, callbacks=[model_checkpoint_callback]
                    )
    return hist


from django import forms

from .models import Feria

class FeriaCreateForm(forms.ModelForm):
	"""docstring for FeriaCreateForm"""
	feria_id = forms.CharField(label = "ID", required=True,
					widget = forms.TextInput(
						attrs={
							"placeholder":"Ingrese la id de la feria",
							"minlength":3
						}
					)
				)
	feria_calle = forms.CharField(label="Nombre calle", required=True,
                widget=forms.TextInput(
                        attrs={
                            "placeholder":"Ingrese el nombre de la calle de la feria",
                            "minlength":5
                        }
                )
            )
	feria_comuna = forms.CharField(label="Comuna feria", required=True,
                widget=forms.TextInput(
                        attrs={
                            "placeholder":"Ingrese el comuna de la feria",
                            "minlength":5
                        }
                	)
                )
	feria_direccion_ini = forms.CharField(label="Direccion inicio feria", required=True,
                widget=forms.TextInput(
                        attrs={
                            "placeholder":"Ingrese la direccion de inicio de la feria",
                            "minlength":5
                        }
                	)
                )
	feria_direccion_fin = forms.CharField(label="Direccion fin feria", required=True,
                widget=forms.TextInput(
                        attrs={
                            "placeholder":"Ingrese la direccion del termino de la feria",
                            "minlength":5
                        }
                	)
            	)
	feria_horario_ini = forms.CharField(label="Horario inicio feria", required=True,
                widget=forms.TextInput(
                        attrs={
                            "placeholder":"Ingrese el horario de inicio de la feria",
                            "minlength":5
                        }
                	)
            	)
	feria_horario_fin = forms.CharField(label="Horario fin feria", required=True,
                widget=forms.TextInput(
                        attrs={
                            "placeholder":"Ingrese el horario de termino de la feria",
                            "minlength":5
                        }
                	)
            	)
	feria_dias = forms.CharField(label="Dias feria", required=True,
                widget=forms.TextInput(
                        attrs={
                            "placeholder":"Ingrese dias de funcionamiento de la feria",
                            "minlength":5
                        }
                	)
            	)
	feria_imagen = forms.ImageField(label="Imagen principal")

	class Meta:
		"""docstring for Meta"""
		model = Feria
		fields = [
			'feria_id',
			'feria_calle',
			'feria_comuna',
			'feria_direccion_ini',
			'feria_direccion_fin',
			'feria_horario_ini',
			'feria_horario_fin',
			'feria_dias',
			'feria_imagen'
		]


#!/usr/bin/env python

'''
Created on Aug 24, 2015

@author: eze
'''

import warnings

from Bio import BiopythonWarning, BiopythonParserWarning

warnings.simplefilter('ignore', BiopythonWarning)
warnings.simplefilter('ignore', BiopythonParserWarning)

import os
import argparse
import multiprocessing
import logging
from tqdm import tqdm

import Bio.SeqIO as  bpio

from SNDG import init_log
from SNDG.Structure.PsiProfile import PsiProfile
from SNDG.Structure.PDBs import PDBs

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument("-fasta", required=True)
    parser.add_argument("-profile_db", "--profile_db", default="/data/uniprot/uniref/uniref90/uniref90.fasta")
    parser.add_argument("-pdbs", "--pdb_seqres", default="/data/pdb/processed/seqs_from_pdb.fasta")
    parser.add_argument("-iter", "--pssm_build_iterations", default=3)
    parser.add_argument("-cpus", "--cpus", default=multiprocessing.cpu_count())  # @UndefinedVariable
    parser.add_argument("-o", "--output_dir", default=None)
    parser.add_argument("-l", "--log_path", default=None)
    parser.add_argument("-e", "--entries", default='/data/pdb/entries.idx')
    parser.add_argument("-d", "--pdb_divided", default="/data/pdb/divided/")
    parser.add_argument("--skipProfile", action='store_true', default=False)
    parser.add_argument("--skipQuality", action='store_true', default=False)

    args = parser.parse_args()

    assert ((not args.skipProfile) or os.path.exists(args.profile_db + ".00.phr"),
            "%s does not exists or it is not indexed" % args.profile_db)
    assert os.path.exists(args.pdb_seqres), "%s does not exists" % args.pdb_seqres
    assert os.path.exists(args.entries), "%s does not exists" % args.entries
    assert os.path.exists(args.pdb_divided), "%s does not exists" % args.pdb_divided
    assert os.path.exists(args.fasta), "%s does not exists" % args.fasta

    if not args.output_dir:
        args.output_dir = os.path.dirname(args.fasta)

    if not args.log_path:
        args.log_path = args.fasta + ".log"
    init_log(args.log_path, logging.INFO)
    _log = logging.getLogger("model_fasta")

    def res_fn(y):
        try:
            res = float(y)
        except:
            res = 30
        return res

    entries = {x.split("\t")[0].lower(): res_fn(x.split("\t")[6]) for x in list(open(args.entries))[3:]}

    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)

    seqs = list(bpio.parse(args.fasta, "fasta"))
    with tqdm(seqs) as pbar:
        for seq in pbar:
            try:
                pbar.set_description(seq.id)
                workdir = os.path.abspath(args.output_dir) + "/" + seq.id + "/"
                if not os.path.exists(workdir):
                    os.makedirs(workdir)
                fasta = workdir + "seq.fasta"
                bpio.write([seq], fasta, "fasta")

                PsiProfile.create_psi_model(seq.id, fasta, args.profile_db,
                                            args.pssm_build_iterations, args.pdb_seqres, workdir,
                                            args.cpus, entries=entries, pdb_divided=args.pdb_divided,
                                            skip_profile=args.skipProfile,skip_quality=args.skipQuality)
            except Exception as ex:
                _log.warning(str(ex))


import hmac
from psycopg2.extensions import quote_ident
from psycopg2.errorcodes import UNDEFINED_TABLE
from django.db import connection, ProgrammingError, InternalError
from django.db.transaction import atomic
from django.utils import timezone

def number_generator(
    tenant_id: str,
    sequence_name: str,
    number_format: str,
    checksum_salt: str,
    checksum_algorithm: str,
    checksum_length: int,
    checksum_format: str,
) -> str:
    if not tenant_id:
        raise ValueError("Tenant-ID not given")

    sequence_name = f'{sequence_name}_{tenant_id}'
    val = None
    i = 0

    while not val and i < 2:
        try:
            with connection.cursor() as cursor:
                with atomic():
                    cursor.execute("SELECT nextval(%s);", (quote_ident(sequence_name, cursor.cursor),))
                    val = cursor.fetchone()
                    break

        except (ProgrammingError, InternalError) as error:
            if error.__cause__.pgcode != UNDEFINED_TABLE:
                raise

            i += 1
            with connection.cursor() as cursor:
                with atomic():
                    cursor.execute("CREATE SEQUENCE IF NOT EXISTS %s START 1;" % quote_ident(sequence_name, cursor.cursor))

    if not val:
        raise ValueError('cannot_obtain_next_number')

    nextnumber = int(val[0])
    now = timezone.now()
    number = number_format % {
        'year': now.year,
        'month': now.month,
        'number': nextnumber,
    }

    if checksum_length > 0:
        checksum_value = hmac.new(checksum_salt, bytes(number, 'utf-8'), digestmod=checksum_algorithm)
        checksum_int = int(checksum_value.hexdigest(), 16)
        checksum_number = checksum_int % 10 ** checksum_length
        checksum = '%0*d' % (checksum_length, checksum_number)
        number = checksum_format % {
            'number': number,
            'checksum': checksum,
        }

    return number


import os
import shutil
from time import *
from datetime import *
from ch.systemsx.cisd.openbis.generic.shared.api.v1.dto import SearchCriteria

MarkerHiSeqComplete = 'RTAComplete.txt'

def createOpenbisTimeStamp(file):
  '''
  Creates a openBIS compatible time stamp of a file time stamp
  '''
  mtime = os.path.getmtime(file)
  lt = localtime(mtime)
  tz = localtime().tm_hour - gmtime().tm_hour
  return (strftime("%Y-%m-%d %H:%M:%S GMT" + "%+.2d" % tz + ":00", lt))

# -----------------------------------------------------------------------------

def extractFlowCellName (runFolderName):
  return runFolderName.split('_')[-1][1:]

# -----------------------------------------------------------------------------

def process(transaction):

  incomingPath = transaction.getIncoming().getAbsolutePath()
  # Get the incoming name
  runFolderName = transaction.getIncoming().getName()
  flowCellName = extractFlowCellName(runFolderName)

  Markerfile = incomingPath + "/" + MarkerHiSeqComplete

  # Search for the sample and check if there is already sample with this name
  search_service = transaction.getSearchService()
  sc = SearchCriteria()
  sc.addMatchClause(SearchCriteria.MatchClause.createAttributeMatch(SearchCriteria.MatchClauseAttribute.CODE, flowCellName));
  foundSamples = search_service.searchForSamples(sc)

  if foundSamples.size() > 0:
      sa = transaction.getSampleForUpdate(foundSamples[0].getSampleIdentifier())
      sa.setPropertyValue("SEQUENCER_FINISHED", createOpenbisTimeStamp(Markerfile))


from django.core.exceptions import ObjectDoesNotExist
from rest_framework import serializers

from hotel_app.models import (
    Booking,
    Hotel,
    City,
)
from hotel_app.utils.common import find_room

class HotelSerializer(serializers.ModelSerializer):
    rooms_beds = serializers.ReadOnlyField(source="get_rooms_beds")

    class Meta:
        model = Hotel
        fields = ("title", "city", "rooms_beds")

class BookingSerializer(serializers.ModelSerializer):
    city = serializers.CharField(read_only=True)
    hotel = serializers.CharField(read_only=True)

    class Meta:
        model = Booking
        fields = ("guest_name", "booking_start", "booking_end", "city", "hotel", "persons")

    @staticmethod
    def get_hotel(city, hotel):
        try:
            city = City.objects.get(title=city)
            hotel = city.hotels.get(title=hotel)
        except ObjectDoesNotExist as e:
            raise serializers.ValidationError(
                {"error": f"{city}:{hotel} - {e}"}
            )
        print(f"Hotel Detected: {hotel}")
        return hotel

    def to_internal_value(self, data):
        print(f"Raw Booking data: {data}")
        self.hotel = self.get_hotel(data.get("city"), data.get("hotel"))
        internal_data = super().to_internal_value(data)
        print(f"Internal Booking data: {internal_data}")
        return internal_data

    def create(self, validated_data):
        print(f"Validated data: {validated_data}")
        booking_start = validated_data.pop("booking_start")
        booking_end = validated_data.pop("booking_end")
        persons = validated_data.pop("persons")
        room = find_room(persons, booking_start, booking_end, self.hotel)
        booking = Booking.objects.create(
            guest_name=validated_data.pop("guest_name"),
            persons=persons,
            booking_start=booking_start,
            booking_end=booking_end,
            room=room,
        )
        print(f"Booking Instance: {booking}")
        return booking


import os
import torchvision
import numpy as np
from PIL import Image
import uuid
from scipy.io import loadmat
from stanford_cars_data_config import CarAugmentedDatasetAttributes
import logging
import shutil
import time
import boto3

AUGMENTED = 'augmented'
IMAGE_DATA_DIR = 'car_ims'
logging.basicConfig(format='%(asctime)s %(message)s', level=logging.INFO)

def clean_up_augmented_images(path_images):
    logging.info("deleting directory: {}".format(path_images))
    shutil.rmtree(path_images)

class StanfordCarsDataAugmentation(object):

    def __init__(self, path_to_data_mat, path_images):
        self.path_images = path_images
        self.path_to_data_mat = path_to_data_mat
        self.data_mat = loadmat(self.path_to_data_mat)[CarAugmentedDatasetAttributes.ANNOTATIONS]
        self.augmented_data_mat = None
        self.rel_path_list = list()
        self.class_list = list()
        self.data_mat_rel_path = list()
        self.data_mat_labels = list()
        self.data_mat_is_augment = list()  # list of booleans indicating if the image is the original or an augment

    def augment_data(self,
                     s3_bucket_name,
                     store_to_disk,
                     store_to_s3,
                     augmentation_multiple,
                     probability,
                     brightness,
                     contrast,
                     saturation,
                     hue
                     ):

        """Augmentation multiple will be used to determine how many times to augment a single image.
        If multiple == 1, vertical flip is favored over horizontal flip.
        If multiple >=2, each image will go through horizontal and vertical flips."""

        if augmentation_multiple < 1:
            raise RuntimeError("Invalid augmentation multiple. Please specify number >=1.")

        # set up directory for augmented data
        augmented_data_directory = os.path.join(self.path_images, IMAGE_DATA_DIR, str(round(time.time())))
        if not os.path.exists(augmented_data_directory):
            logging.info("making directory for augmented data: {}".format(augmented_data_directory))
            os.makedirs(augmented_data_directory)

        s3_client = boto3.resource('s3')

        for index in range(self.data_mat.shape[1]):

            transformations = self.generate_transformation_color_jitter(augmentation_multiple=augmentation_multiple,
                                                                        probability=probability,
                                                                        brightness=brightness,
                                                                     contrast=contrast,
                                                                        saturation=saturation,
                                                                        hue=hue)

            # store data information in unnested list form
            self.unnest_data_mat(index)

            # apply transformations to image
            transformed_images_list = self.apply_transformations(index, transformations)

            if store_to_s3:
                for transformed_image in transformed_images_list:
                    # store to disk
                    transformed_image_path = self.save_transform_to_disk(index, augmented_data_directory,
                                                                       transformed_image)
                    # store to s3
                    self.save_transform_to_s3(s3_client, s3_bucket_name, transformed_image_path)
            elif store_to_disk:
                for transformed_image in transformed_images_list:
                    # store to disk
                    self.save_transform_to_disk(index, augmented_data_directory,
                                                                       transformed_image)
            else:
                logging.debug("no action specified for storing images. might be a problem for the future")

        logging.info("total number of images in directory post augmentation: {}".format(len(
            os.listdir(augmented_data_directory))))

        return augmented_data_directory

    def generate_transformation_color_jitter(self,
                                augmentation_multiple,
                                probability,
                                brightness,
                                contrast,
                                saturation,
                                hue
                                ):

        """Function will create transformations to horizontally flip and jitter color.
         A single augmentation will involve both transformations being performed."""
        transformations = list()
        for multiple in range(0, augmentation_multiple):
            horizontal_flip_transform = torchvision.transforms.RandomHorizontalFlip(probability)
            color_transform = torchvision.transforms.ColorJitter(brightness=brightness,
                                                                 contrast=contrast,
                                                                 saturation=saturation,
                                                                 hue=hue)
            transformations.append([horizontal_flip_transform, color_transform])
        return transformations

    def apply_transformations(self, index, transformations):
        """for current image, apply all transformations, save to disk, append information to data structure"""

        transformed_images_list = list()
        for transformation_group in transformations:
            composed_transform = torchvision.transforms.Compose(transformation_group)
            logging.debug("current transform: %s", composed_transform)
            image_path = os.path.join(self.path_images, self.data_mat[0][index][
                CarAugmentedDatasetAttributes.REL_IMAGE_PATH][0])
            logging.info("transforming image: {}".format(image_path))
            image = Image.open(image_path)
            transformed_image = composed_transform(image)
            transformed_images_list.append(transformed_image)

        return transformed_images_list

    def save_transform_to_disk(self, index, augmented_data_directory, transformed_image):
        augmented_image_epithet = AUGMENTED
        augmented_image_filename = '{}_{}.jpg'.format(uuid.uuid4().hex, augmented_image_epithet)
        augmented_image_filepath = os.path.join(augmented_data_directory, augmented_image_filename)

        # adding to augmented data matrix
        self.rel_path_list.append(augmented_image_filepath)
        self.class_list.append(self.data_mat[0][index][CarAugmentedDatasetAttributes.LABEL][0][0])

        # saving to disk
        transformed_image_path = os.path.join(self.path_images, augmented_image_filepath)
        logging.info("saving transformed image to: %s", transformed_image_path)

        transformed_image.save(transformed_image_path)
        return transformed_image_path

    def save_transform_to_s3(self, s3_client, s3_bucket_name, transformed_image_path):
        image_opened_s3 = open(transformed_image_path, 'rb')
        s3_client.Bucket(s3_bucket_name).put_object(Key=transformed_image_path, Body=image_opened_s3)
        image_opened_s3.close()

    def unnest_data_mat(self, index):
        """ Store matrix data as unnested lists. Converts relative path of image location to absolute path."""
        self.data_mat_rel_path.append(os.path.join(self.path_images, self.data_mat[0][index][
                                          CarAugmentedDatasetAttributes.REL_IMAGE_PATH][0]))
        self.data_mat_labels.append(self.data_mat[0][index][CarAugmentedDatasetAttributes.LABEL][0][0])
        self.data_mat_is_augment.append(False)

    def get_all_data_matrix(self):
        self.data_mat_rel_path.extend(self.rel_path_list)
        self.data_mat_labels.extend(self.class_list)
        self.data_mat_is_augment.extend([True] * len(self.rel_path_list))

        all_data_mat = np.empty(len(self.data_mat_rel_path), dtype={'names': (
            CarAugmentedDatasetAttributes.REL_IMAGE_PATH, CarAugmentedDatasetAttributes.LABEL,
            CarAugmentedDatasetAttributes.IS_AUGMENT), 'formats': ('<U200', 'uint8', '?')})
        all_data_mat[CarAugmentedDatasetAttributes.REL_IMAGE_PATH] = self.data_mat_rel_path
        all_data_mat[CarAugmentedDatasetAttributes.LABEL] = self.data_mat_labels
        all_data_mat[CarAugmentedDatasetAttributes.IS_AUGMENT] = self.data_mat_is_augment

        logging.info("total number of images for training: {}".format(len(all_data_mat)))
        return all_data_mat


# -*- coding: utf-8 -*-
"""
Created on Thu Aug 27 09:12:01 2020

@author: ler45
"""

# Function for calculating global average temperature (not smoothed), relative to
# the pre-industrial control period.
# This version for CMIP5 data

def gmst_calc(model, RCP, Yend, rcpEns):

    import numpy as np
    from matplotlib import pyplot as plt
    import os
    import xarray as xr
    from statsmodels.nonparametric.smoothers_lowess import lowess
    import cftime

    #This block of code creates a list that contains all of the files we need to analyse
    #a single simulation run. `historical` simulations cover the time period 1850-2015
    #and scenario (e.g. `ssp245`) simulations cover the time period 2015-2100. So
    #a single simulation from 1850-2100 needs both the `historical` and the predictive
    #`ssp245` to get full temporal coverage.
    #
    #The `DATAPATH` is the path to the directory on your computer that contains
    #the cmip6 data.

    ###############################
    ####----USER-CHANGEABLE----####

    # Set the parameters of interest:
    variable_id = 'tas'
    table_id = 'Amon'

    scenarios = ['historical',RCP]

    DATAPATH = '/Volumes/GeoCCRI_01/users-data/douglahu/CMIP5_data/'+variable_id+'/'+table_id+'/'+model+'/'

    # Set whether or not to use the use_cftime option in xr.open_dataset
    # (this avoids issues if the dates go past 2262, but might slow performance. Not sure.
    usingCFtime = True

    ####----END USER-CHANGEABLE----####
    ###################################

    os.chdir(DATAPATH)

    ensembles = []
    fileList = []

    print('Checkpoint1')

    for scenario in scenarios:

        if scenario == 'historical':
            # Get the ensemble member for the historical run
            ensembleList = np.sort(os.listdir(DATAPATH+scenario))
            print(ensembleList)
            if ensembleList[0] == '.DS_Store':
                ensembleList = ensembleList[1:]
            if len(ensembleList)  > 1:
                print('WARNING: Multiple ensembles exist for this instance. Only using first.')
            ensembles.append(ensembleList[0])
        else:
            ensembles.append(rcpEns) #Use the specified variant
            # This assumes that ssp585 has all the same variants as ssp534-over

        # Pick the ensemble relating to this scenario (this can vary between scenarios)
        ensemble = ensembles[scenarios.index(scenario)]

        # Create a file list of all raw data to collate
        try:
            files = [f for f in os.listdir(DATAPATH+scenario+'/'+ensemble) if ensemble in f]
        except:
            print('Error in reading specified filename. Might be misalignment between ssp585 and ssp534-over.')
        for f in files:
            fileList.append(DATAPATH+scenario+'/'+ensemble+'/'+f)

    fileList.sort()
    #print(fileList)

    #Uncomment lines below if data are across multiple pressure levels.

    def preprocess(ds):
            #data = ds.isel(plev=slice(None,7))
            #data = data.bfill('plev')
            #data = data.isel(plev=0)

            data = ds.groupby('time.year').mean('time').load()
            try:
                test=len(data.lat)
            except:
                print('Renaming latitude and longitude')
                data=data.rename({'latitude':'lat','longitude':'lon'})

            return data

    print('Checkpoint2')

    # Caution: some of the historical and SSP files overlap in time, so open_mfdataset fails
    # So, I set up a less efficient, but still functional alternative
    try:
        with xr.open_mfdataset(fileList, preprocess=preprocess, parallel=True, use_cftime=usingCFtime) as file:
            data = file
        # EC-Earth3 has a single digit off in the 13th decimal place for two
        # of the latitude values. This ruins the combined dataset. Need to
        # check this and throw an error to prompt the exception code to run.
        print('Checking lengths')
        init_file = fileList[0]
        initData = xr.open_dataset(init_file, use_cftime=True)
        try:
            test=len(initData.lat)
        except:
            print('Renaming latitude and longitude')
            initData=initData.rename({'latitude':'lat','longitude':'lon'})
        if len(data.lat) != len(initData.lat):
            fileList.index('boguscheck') # Intentially throw an error
        # Also throw an error if there are any repeated years or gaps in the data
        if int(max(data.year).values)-int(min(data.year).values)+1 != len(data.year):
            print('Something is off with the years. Applying workaround.')
            fileList.index('boguscheck') # Intentially throw an error
    except:
        print('xr.open_mfdataset() failed. Applying workaround.')
        init_file = fileList[0]

        # Initiate the xarray that will have all the data
        allData = xr.open_dataset(init_file, use_cftime=True)
        try:
            test=len(allData.lat)
        except:
            print('Renaming latitude and longitude')
            allData=allData.rename({'latitude':'lat','longitude':'lon'})
        realLats = allData.lat

        for i in fileList:
            # Skip over the initial file
            if i != init_file:
                print(i)
                # Open and process this file
                dset = xr.open_dataset(i, use_cftime=True)
                try:
                    test=len(dset.lat)
                except:
                    print('Renaming latitude and longitude')
                    dset=dset.rename({'latitude':'lat','longitude':'lon'})
                # Check for and correct latitude mismatch
                if np.sum(abs(dset.lat.values-realLats.values)) != 0.0:
                    print('Latitude mismatch = ',str(np.sum(abs(dset.lat.values-realLats.values))))
                    dset = dset.assign_coords(lat=realLats)

                maxTime = np.max(allData.time.values)
                minTime = np.min(dset.time.values)

                if maxTime > minTime:
                    print('Overlap between historical & simulation. Truncating sim data.')
                    if maxTime.month == 12:
                        startRCP = str(maxTime.year+1)+'-01'
                    else:
                        startRCP = str(maxTime.year)+'-'+str(maxTime.month+1)
                    endRCP = str(np.max(dset.time.values).year)+'-'+str(np.max(dset.time.values).month)
                    dset = dset.loc[dict(time=slice(startRCP, endRCP))]

                # Concatenate along time axis
                allData = xr.concat((allData, dset), dim="time")

        print('Checkpoint2.2')

        try:
            data = allData.groupby('time.year').mean('time')
        except:
            print('Sorting failed. Calendar mismatch. Engaging workaround.')
            realDates = []
            for j in range(len(allData.time)):
                thisDate = allData.time.values[j]
                realDates.append(cftime.datetime(thisDate.year,thisDate.month,thisDate.day))
            allData = allData.assign_coords(time=realDates)
            print('Checkpoint 2.3')
            data = allData.groupby('time.year').mean('time')

    print(data.lat.shape)

    # Before fitting the polynomial, it's beneficial to truncate the data after
    # 2100. Inspecting some of the fits showed that having extended datasets really
    # skews the fit over important periods.
    data = data.loc[dict(year=slice(1850,Yend))]

    # calculate tasAverage (note that GMST has been used to refer to the smoothed
    # curve in other scripts.)

    weights = np.cos(np.deg2rad(data['lat']))

    data['tasAverage'] = (data[variable_id]*weights).sum('lat').mean('lon')/np.sum(weights)

    # Now load the piControl data to calculate the baseline to subtract away,
    # leaving tasAverage as degrees K above pre-industrial.

    # Get the ensemble and grid label
    ensembleList = np.sort(os.listdir(DATAPATH+'piControl/'))
    print(ensembleList)
    if ensembleList[0] == '.DS_Store':
        ensembleList = ensembleList[1:]
    if len(ensembleList)  > 1:
        print('WARNING: Multiple ensembles exist for this instance. Only using first.')
    ensembles.append(ensembleList[0])
    # Pick the ensemble relating to piControl
    ensemble = ensembles[-1]

    print('Checkpoint3')

    fileList = []
    files = [f for f in os.listdir(DATAPATH+'piControl/'+ensemble) if ensemble in f]
    for f in files:
        fileList.append(DATAPATH+'piControl/'+ensemble+'/'+f)

    fileList.sort()
    #print(fileList)

    print('Checkpoint3.1')

    # HadGEM2-ES has this annoying thing where each file ends in November and
    # the next one starts in December, so the pre-processing doesn't work.
    try:
        with xr.open_mfdataset(fileList, preprocess=preprocess, parallel=True, use_cftime=usingCFtime) as file:
            piCdata = file
    except:
        print('xr.open_mfdataset() failed. Applying workaround.')
        init_file = fileList[0]

        # Initiate the xarray that will have all the data
        allData = xr.open_dataset(init_file, use_cftime=True)
        try:
            test=len(allData.lat)
        except:
            print('Renaming latitude and longitude')
            allData=allData.rename({'latitude':'lat','longitude':'lon'})
        realLats = allData.lat

        for i in fileList:
            # Skip over the initial file
            if i != init_file:
                print(i)
                # Open and process this file
                dset = xr.open_dataset(i, use_cftime=True)
                try:
                    test=len(dset.lat)
                except:
                    print('Renaming latitude and longitude')
                    dset=dset.rename({'latitude':'lat','longitude':'lon'})
                # Check for and correct latitude mismatch
                if np.sum(abs(dset.lat.values-realLats.values)) != 0.0:
                    print('Latitude mismatch = ',str(np.sum(abs(dset.lat.values-realLats.values))))
                    dset = dset.assign_coords(lat=realLats)

                # Concatenate along time axis
                allData = xr.concat((allData, dset), dim="time")

        print('Checkpoint3.1b')

        try:
            piCdata = allData.groupby('time.year').mean('time')
        except:
            print('Sorting failed. Calendar mismatch. Engaging workaround.')
            realDates = []
            for j in range(len(allData.time)):
                thisDate = allData.time.values[j]
                realDates.append(cftime.datetime(thisDate.year,thisDate.month,thisDate.day))
            allData = allData.assign_coords(time=realDates)
            print('Checkpoint 3.1c')
            piCdata = allData.groupby('time.year').mean('time')

    # Only use the last 200 years of the piControl data, to allow for ramp-up
    piCend = np.max(piCdata.year.values)
    print('Checkpoint3.2')
    piCdata = piCdata.loc[dict(year=slice(piCend-199,piCend))]
    print('Checkpoint3.3')
    # Area-average gmst over the piControl period:
    weights = np.cos(np.deg2rad(piCdata['lat']))
    print('Checkpoint3.4')
    piCgmstAvg = (piCdata*weights).sum('lat').mean('lon')/np.sum(weights)
    print('Checkpoint3.5')
    piCgmstAvg = piCgmstAvg[variable_id].mean('year')
    print(piCgmstAvg)
    print('Checkpoint3.6')

    # Subtract this from tasAverage to get GMST relative to piC baseline

    # But MIROC5 has a weird issue where the piC data are way too cold.
    # In that case only, I'll calculate GMST relative to 1850-1900 avg temps.
    if model=='MIROC5':
        histGmstAvg = data['tasAverage'].loc[dict(year=slice(1850,1900))].mean('year')
        data['tasAverageRel'] = data['tasAverage'] - histGmstAvg
    else:
        data['tasAverageRel'] = data['tasAverage'] - piCgmstAvg

    return data


import random
import sys

from ckiptagger import data_utils, construct_dictionary, WS, POS, NER
#data_utils.download_data_gdown("./")
#data_utils.download_data_url("./") # iis-ckip
#data_utils.download_data_gdown("./") # gdrive-ckip

'''
Usage:
python split.py src_fpath tgt_fpath new_data_dir
'''

def split(src_fpath, tgt_fpath, nsrc='zh', ntgt='tl', ratio=(0.9, 0.05, 0.05), new_data_dir=''):
  src_fp = open(src_fpath, encoding='utf-8')
  tgt_fp = open(tgt_fpath, encoding='utf-8')

  src_train, src_test, src_val = open(new_data_dir + 'train.' + nsrc, 'w', encoding='utf-8'), \
    open(new_data_dir + 'test.' + nsrc, 'w', encoding='utf-8'), open(new_data_dir + 'valid.' + nsrc, 'w', encoding='utf-8')
  tgt_train, tgt_test, tgt_val = open(new_data_dir + 'train.' + ntgt, 'w', encoding='utf-8'), \
    open(new_data_dir + 'test.' + ntgt, 'w', encoding='utf-8'), open(new_data_dir + 'valid.' + ntgt, 'w', encoding='utf-8')

  src, tgt = src_fp.readlines(), tgt_fp.readlines()
  #print(src)
  #print(len(src))
  #p = input('')
  #print(tgt)
  #print(len(tgt))
  #p = input('')
  colline=1
  ws = WS("./data")
  for s, t in zip(src, tgt):
      rand = random.random()
      if 0 < rand <= ratio[0]:
        #if len(t) < 6:
        #    p=input("")
        #    continue
        #if len(s) < 6:
        #    p=input("")
        #    continue
        _ws = ws([t])
        #print(_ws)
        t=" ".join(_ws[0])
        print(len(s),s.lower())
        print(len(t),t)
        if len(t) < 6:
            p=input("")
            continue
        if len(s) < 6:
            p=input("")
            continue
        #p=input("")
        src_train.write(s.lower())
        tgt_train.write(t)
        print("---- train"+"-"*100)
      elif ratio[0] < rand <= ratio[0] + ratio[1]:
        if len(t) == 1 or len(t) < 6:
            continue
        if len(s) == 1 or len(s) < 6:
            continue
        _ws = ws([t])
        #print(_ws)
        t=" ".join(_ws[0])
        print(len(s),s.lower())
        print(len(t),t)
        print("---- test"+"-"*100)
        src_test.write(s.lower())
        tgt_test.write(t)
      else:
        if len(t) == 1 or len(t) < 6:
            continue
        if len(s) == 1 or len(s) < 6:
            continue
        _ws = ws([t])
        #print(_ws)
        t=" ".join(_ws[0])
        print(len(s),s.lower())
        print(len(t),t)
        print("---- val"+"-"*100)
        src_val.write(s.lower())
        tgt_val.write(t)

  src_fp.close()
  tgt_fp.close()
  src_train.close()
  src_test.close()
  src_val.close()
  tgt_train.close()
  tgt_test.close()
  tgt_val.close()

if __name__ == '__main__':
    split(src_fpath=sys.argv[1], tgt_fpath=sys.argv[2], nsrc='tl', ntgt='zh', ratio=(0.8, 0.1, 0.1), new_data_dir=sys.argv[3])


from typing import List

class Codec:
    def encode(self, strs: List[str]) -> str:
        """Encodes a list of strings to a single string.
        """
        encoded = ""
        for item in strs:
            # eg. "5#Hello5#World"
            encoded += str(len(item)) + "#" + item
        return encoded

    def decode(self, s: str) -> List[str]:
        """Decodes a single string to a list of strings.
        """
        index_s = 0
        decoded = []
        while index_s < len(s):
            index_num = index_s
            # Find out how long is the string by getting a number
            # comes before #, and a number can be more than one digit
            while s[index_num] != "#":
                index_num += 1
            # eg. "5#Hello5#World" index num would be 1, 8
            # [0:1] [7:8]
            length_item = int(str(s[index_s:index_num]))
            # +1 to indexes to skip # added
            decoded.append(s[index_num+1:index_num+1+length_item])
            index_s = index_num + 1 + length_item
        return decoded

strs = ["Hello","World"]
codec = Codec()
print("Result of encode is", codec.encode(strs))
print("Result of encode -> decode is", codec.decode(codec.encode(strs)))


'''
Task 34
Seth Maurice-Brant
'''

checkNum = input("Enter number to check beauty; ")

prevDigit = ""
for digit in checkNum:
    if int(digit) == int(prevDigit)+1:
        prevDigit = digit
    else:
        print("Not beatiful")
        break


import numpy as np
from math import *
import random

l1=50
l2=20
l3=120
l4=155

def legIK(point):
        (x,y,z)=(point[0],point[1],point[2])
        (_l1,_l2,_l3,_l4)=(l1,l2,l3,l4)
        try:
            F=sqrt(x**2+y**2-_l1**2)
        except ValueError:
            print("Error in legIK with x {} y {} and _l1 {}".format(x,y,_l1))
            F=_l1
        G=F-_l2
        H=sqrt(G**2+z**2)
        print(H)
        theta1=-atan2(y,x)-atan2(F,-_l1)
        print(theta1)
        D=(H**2-_l3**2-_l4**2)/(2*_l3*_l4)
        print(D)
        try:
            theta3=acos(D)
        except ValueError:
            print("Error in legIK with x {} y {} and D {}".format(x,y,D))
            theta3=0
        theta2=atan2(z,G)-atan2(_l4*sin(theta3),_l3+_l4*cos(theta3))

        return(theta1,theta2,theta3)

a = np.array([1200, 1200, 1200])
b = legIK(a)

print(b)


import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torch.autograd import Variable
import itertools
from torch import FloatTensor as FT

try:
    from utils import *
except:
    from .utils import *

# class Cholesky(torch.autograd.Function):
#     def forward(ctx, a):
#         l = torch.cholesky(a, False)
#         ctx.save_for_backward(l)
#         return l
#
#     def backward(ctx, grad_output):
#         l, = ctx.saved_variables
#         linv = l.inverse()
#         inner = torch.tril(torch.mm(l.t(), grad_output)) * torch.tril(
#             1.0 - Variable(l.data.new(l.size(1)).fill_(0.5).diag()))
#         s = torch.mm(linv.t(), torch.mm(inner, linv))
#         return s

class AE_encoder(nn.Module):
    def __init__(
            self,
            device,
            structure_config
    ):

        super(AE_encoder, self).__init__()
        self.device = device
        self.structure_config = structure_config
        activation = structure_config['encoder_layers']['activation']
        layer_dims = structure_config['encoder_layers']['layer_dims']

        # ================
        # Concatenate the input
        # ================

        self.discrete_column_dims = structure_config['discrete_column_dims']
        self.num_discrete_columns = structure_config['num_discrete']
        self.num_real_columns = structure_config['num_real']

        input_dim = self.num_discrete_columns + self.num_real_columns
        layers = []
        # Pairs <  output_dim, activation >
        num_layers = len(layer_dims)
        for idx in range(num_layers):
            op_dim = layer_dims[idx]
            layers.append(nn.Linear(
                input_dim, op_dim
            ))
            if idx == num_layers - 1: activation = 'none'
            if activation == 'tanh':
                layers.append(
                    nn.Tanh()
                )
            elif activation == 'relu':
                layers.append(
                    nn.ReLU()
                )
            elif activation == 'sigmoid':
                layers.append(
                    nn.Sigmoid()
                )
            elif activation == 'none':
                pass

            input_dim = op_dim

        self.FC_z = nn.Sequential(*layers)
        return

    def forward(self, X):

        # real_x_0 = X[:, -self.num_real_columns:].type(torch.FloatTensor).to(self.device)
        # discrete_x_0 = X[:, :self.num_discrete_columns].type(torch.LongTensor).to(self.device)
        # discrete_x_1 = torch.chunk(discrete_x_0, self.num_discrete_columns, dim=1)
        # # discrete_x_1 is an array
        # res = []
        # column_name_list = list(self.discrete_column_dims.keys())
        # for idx in range(self.num_discrete_columns):
        #     col = column_name_list[idx]
        #     _x = discrete_x_1[idx].to(self.device)
        #     n_cat = self.discrete_column_dims[col]
        #     _x = F.one_hot(_x, n_cat).type(FT).squeeze(1).to(self.device)
        #     res.append(_x)
        #
        # res.append(real_x_0)
        # x_concat = torch.cat(res, dim=1)

        # ==============
        # Conactenated input : res
        # ==============
        op = self.FC_z(X)
        return op

class AE_decoder(nn.Module):
    def __init__(
            self,
            device,
            structure_config=None,
    ):
        super(AE_decoder, self).__init__()
        self.device = device
        self.structure_config = structure_config
        self.discrete_column_dims = structure_config['discrete_column_dims']
        self.num_discrete_columns = structure_config['num_discrete']
        self.num_real_columns = structure_config['num_real']
        final_output_dim = structure_config['final_output_dim']
        activation = structure_config['decoder_layers']['activation']
        layer_dims = structure_config['decoder_layers']['layer_dims']

        layers = []
        inp_dim = layer_dims[0]
        num_layers = len(layer_dims)
        for idx in range(1, num_layers):
            op_dim = layer_dims[idx]
            layers.append(nn.Linear(inp_dim, op_dim))
            if activation == 'tanh':
                layers.append(
                    nn.Tanh()
                )
            elif activation == 'sigmoid':
                layers.append(
                    nn.Sigmoid()
                )
            elif activation == 'relu':
                layers.append(
                    nn.ReLU()
                )
            elif activation == 'none':
                layers.append(
                    nn.Sigmoid()
                )
            inp_dim = op_dim

        self.FC_z = nn.Sequential(*layers)
        return

    def forward(self, z):
        res = self.FC_z(z)
        return res

class AE(nn.Module):

    def __init__(
            self,
            device,
            encoder_structure_config,
            decoder_structure_config
    ):
        super(AE, self).__init__()
        self.device = device
        self.encoder = AE_encoder(
            device,
            encoder_structure_config
        )
        self.decoder = AE_decoder(
            device,
            decoder_structure_config
        )

        self.encoder.to(self.device)
        self.decoder.to(self.device)
        self.mode = None
        return

    def forward(self, x):
        z = self.encoder(x)
        if self.mode == 'compress':
            return z
        x_recon = self.decoder(z)
        return x_recon, z

class DaGMM(nn.Module):

    def __init__(
            self,
            device,
            encoder_structure_config,
            decoder_structure_config,
            n_gmm=2,
            ae_latent_dim=1,
            fc_dropout=0.5
    ):

        super(DaGMM, self).__init__()
        self.device = device
        self.fc_dropout = fc_dropout
        self.encoder = AE_encoder(
            device,
            encoder_structure_config
        )
        self.encoder = self.encoder.to(self.device)
        self.decoder = AE_decoder(
            device,
            decoder_structure_config
        )

        self.decoder = self.decoder.to(self.device)
        latent_dim = ae_latent_dim + 2
        layers = []
        layers += [nn.Linear(latent_dim, 10)]
        layers += [nn.Tanh()]
        layers += [nn.Dropout(p=self.fc_dropout)]
        layers += [nn.Linear(10, n_gmm)]
        layers += [nn.Softmax(dim=1)]

        self.estimation = nn.Sequential(*layers)

        self.register_buffer("phi", torch.zeros(n_gmm))
        self.register_buffer("mu", torch.zeros(n_gmm, latent_dim))
        self.register_buffer("cov", torch.zeros(n_gmm, latent_dim, latent_dim))

    def relative_euclidean_distance(self, a, b):
        return (a - b).norm(2, dim=1) / a.norm(2, dim=1)

    def forward(self, x):
        enc = self.encoder(x)
        dec = self.decoder(enc)
        rec_cosine = F.cosine_similarity(x, dec, dim=1)
        rec_euclidean = self.relative_euclidean_distance(x, dec)
        z = torch.cat([enc, rec_euclidean.unsqueeze(-1), rec_cosine.unsqueeze(-1)], dim=1)
        gamma = self.estimation(z)
        return enc, dec, z, gamma

    def compute_gmm_params(self, z, gamma):
        N = gamma.size(0)
        # K
        sum_gamma = torch.sum(gamma, dim=0)

        # K
        phi = (sum_gamma / N)

        self.phi = phi.data

        # K x D
        mu = torch.sum(gamma.unsqueeze(-1) * z.unsqueeze(1), dim=0) / sum_gamma.unsqueeze(-1)
        self.mu = mu.data
        # z = N x D
        # mu = K x D
        # gamma N x K

        # z_mu = N x K x D
        z_mu = (z.unsqueeze(1) - mu.unsqueeze(0))

        # z_mu_outer = N x K x D x D
        z_mu_outer = z_mu.unsqueeze(-1) * z_mu.unsqueeze(-2)

        # K x D x D
        cov = torch.sum(gamma.unsqueeze(-1).unsqueeze(-1) * z_mu_outer, dim=0) / sum_gamma.unsqueeze(-1).unsqueeze(-1)
        self.cov = cov.data

        return phi, mu, cov

    def compute_energy(self, z, phi=None, mu=None, cov=None, size_average=True):
        if phi is None:
            phi = to_var(self.phi)
        if mu is None:
            mu = to_var(self.mu)
        if cov is None:
            cov = to_var(self.cov)

        k, D, _ = cov.size()
        z_mu = (z.unsqueeze(1) - mu.unsqueeze(0))
        cov_inverse = []
        det_cov = []
        cov_diag = 0
        eps = 1e-12
        for i in range(k):
            # K x D x D
            cov_k = cov[i] + to_var(torch.eye(D) * eps)
            cov_inverse.append(torch.inverse(cov_k).unsqueeze(0))
            tmp = np.linalg.det(cov_k.data.cpu().numpy() * (2 * np.pi))
            det_cov.append(tmp)
            cov_diag = cov_diag + torch.sum(1 / cov_k.diag())

        # K x D x D
        cov_inverse = torch.cat(cov_inverse, dim=0)
        # K
        det_cov = to_var(torch.from_numpy(np.float32(np.array(det_cov))))

        # N x K
        exp_term_tmp = -0.5 * torch.sum(
            torch.sum(
                z_mu.unsqueeze(-1) * cov_inverse.unsqueeze(0), dim=-2
            ) * z_mu,
            dim=-1
        )
        # for stability (logsumexp)
        k = (exp_term_tmp).clamp(min=0)

        max_val = torch.max(k, dim=1, keepdim=True)[0]
        exp_term = torch.exp(exp_term_tmp - max_val)
        sample_energy = -max_val.squeeze() - torch.log(
            torch.sum(phi.unsqueeze(0) * exp_term / (torch.sqrt(det_cov)).unsqueeze(0), dim=1) + eps)

        if size_average:
            sample_energy = torch.mean(sample_energy)

        return sample_energy, cov_diag

    def loss_function(self, x, x_hat, z, gamma, lambda_energy, lambda_cov_diag):
        recon_error = torch.mean((x - x_hat) ** 2)
        phi, mu, cov = self.compute_gmm_params(z, gamma)
        sample_energy, cov_diag = self.compute_energy(z, phi, mu, cov)
        loss = recon_error + lambda_energy * sample_energy + lambda_cov_diag * cov_diag

        return loss, sample_energy, recon_error, cov_diag


#!/usr/bin/env python
import os
import os.path as osp
import shutil
import sys
import tempfile
import warnings
from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser, Namespace
from dataclasses import dataclass
from typing import Tuple

from esgpull import Esgpull, Query

from cmip6_utils.file import download_file
from cmip6_utils.misc import BC, verify_checksum

@dataclass
class ESGFFile:
    filename: str
    url: str
    checksum: str
    local_path: str
    data_node: str

def cmip_path_to_query(path: str) -> Query:
    if not path.startswith("CMIP6"):
        raise RuntimeError("path should be described with root directory CMIP6")

    path = path.strip().split("/")

    query = Query()
    query.selection.project = path[0]
    query.selection.activity_id = path[1]
    query.selection.institution_id = path[2]
    query.selection.source_id = path[3]
    query.selection.experiment_id = path[4]
    query.selection.variant_label = path[5]
    query.selection.table_id = path[6]
    query.selection.variable_id = path[7]
    query.selection.grid_label = path[8]

    return query

def cli() -> Namespace:
    parser = ArgumentParser(
        formatter_class=ArgumentDefaultsHelpFormatter,
        description=(
            "This script can be used to download all available files in a dataset. "
            "For example for /the/full/local/path/to/dataset, the script will query all "
            "all available files and download the missing files to that location."
        ),
    )
    parser.add_argument("dataset", type=str, help="Full local path to the dataset to download")
    args = parser.parse_args(args=None if sys.argv[1:] else ["--help"])
    args.dataset = args.dataset.rstrip("/")

    if args.dataset.endswith(".nc"):
        raise ValueError("The supplied argument should be the path to the datset, not a specific file")

    return args

def split_dataset_path(path: str) -> Tuple[str, str]:
    i = path.find("CMIP6")
    return path[:i].rstrip("/"), path[i:].strip("/")

def dataset_version(path: str) -> str:
    version = path.strip().split("/")[-1]
    assert version.startswith("v")
    return version

def print_query(query: Query):
    print(BC.bold("Query selection:"))
    for facet in query.selection._facets:
        print(f"{facet.name:15s}: {facet.value}")
    print("\n")

def main():
    args = cli()
    find_download_missing_files(args.dataset)

def find_download_missing_files(dataset: str):
    # args = cli()

    path_to_cmip6_data, cmip6_dataset_structure = split_dataset_path(dataset)
    version = dataset_version(dataset)

    # print(f"{BC.bold('Dataset path:')} {args.dataset}")
    print(f"{BC.bold('Local path to CMIP6 data:')} {path_to_cmip6_data}")
    print(f"{BC.bold('CMIP6 dataset structure :')} {cmip6_dataset_structure}")
    print(f"{BC.bold('Dataset version         :')} {version}")
    print("\n")

    query = cmip_path_to_query(cmip6_dataset_structure)

    print_query(query)

    query.options.distrib = True  # default=False
    query.options.replica = True

    esg = Esgpull(path="/home/dchandan/esgpull_profiles/scratch")
    search_results = esg.context.files(query, max_hits=None)

    dataset_files = {}

    print(f"{BC.warn('# of files found from search: ')}{len(search_results)}")

    for res in search_results:
        if version == res.version:
            filename = res.filename
            data_ = (filename, res.url, res.checksum, res.local_path, res.data_node)
            if filename in dataset_files:
                dataset_files[filename].append(ESGFFile(*data_))
            else:
                dataset_files[filename] = [ESGFFile(*data_)]

    print(f"{BC.warn('# of unique files: ')}{len(dataset_files)}")

    changes_made = False
    for filename, entries in dataset_files.items():
        if not osp.exists(osp.join(dataset, filename)):
            success = False
            print(f"Found missing file: {filename}")

            tf_ = tempfile.NamedTemporaryFile()
            for item in entries:
                print(f" ---> Attempting download from: {item.data_node}")
                err = download_file(item.url, tf_.name)

                if err == 200:
                    print(f" ---> Download {BC.okgreen('successful')}")
                    vrfy = verify_checksum(tf_.name, item.checksum)
                    if vrfy:
                        print(f" ---> Checksum {BC.okgreen('PASS')}")
                        success = True
                        localpath = osp.join(path_to_cmip6_data, item.local_path)
                        print(f" ---> Copying downloaded file to {localpath}")
                        os.makedirs(localpath, exist_ok=True)
                        local_filename = osp.join(localpath, filename)
                        shutil.copy2(tf_.name, local_filename)
                        vrfy2 = verify_checksum(local_filename, item.checksum)
                        if not vrfy2:
                            print(f" ---> {BC.fail('Checksum of copied file failed')}")
                            os.remove(local_filename)
                            continue

                        changes_made = True
                        break
                    else:
                        print(f" ---> Checksum {BC.fail('FAIL')}")
                else:
                    print(f" ---> Download {BC.fail('unsuccessful')}")

            if not success:
                print(f" ---> {BC.fail('Unable to download this file from any source')}")

    if not changes_made:
        print(f"{BC.warn('No missing files were found')}")

if __name__ == "__main__":
    main()


# written by Dustin Frandle 2/22/2021
# designed to be used on a NumWorks Calculator
import statisticaltdm
import synchronoustdm
import pulsestuffing

print("Select TDM")
print("1. Synchronous/Multilevel/Multiple-slot")
print("2. Statistical")
print("3. Pulse-Stuffing")
entry = input("")
if entry == "1":
    a = input("smallest input channel in bits/s: ")
    b = input("largest input channel in bits/s: ")
    c = input("the size of the time slot: ")
    d = input("additional bits: ")
    e = input("sum of the data rate of each channel: ")
    print("---------------------------------------")
    synchronoustdm.calculateSyncTDM(float(a), float(b), float(c), float(d), float(e))
elif entry == "2":
    a = input("number of slots: ")
    b = input("the size in bits of a character: ")
    c = input("the data rate of incoming character per second: ")
    d = input("the size of the time slot: ")
    e = input("number of sources: ")
    print("---------------------------------------")
    statisticaltdm.calculateStatTDM(float(a), float(b), float(c), float(d), float(e))
elif entry == "3":
    a = input("smallest input channel in bits/s: ")
    b = input("largest input channel in bits/s: ")
    c = input("the size of the time slot: ")
    d = input("additional bits: ")
    e = input("number of input channels: ")
    print("---------------------------------------")
    pulsestuffing.calculateStuffTDM(float(a), float(b), float(c), float(d), float(e))


from .code import Code, Flag
from .core import CodeTransformer
from . patterns import (
    matchany,
    not_,
    option,
    or_,
    pattern,
    plus,
    seq,
    var,
)
from . import instructions
from . import transformers
from .utils.pretty import a, d, display, pprint_ast, pformat_ast
from ._version import get_versions

__version__ = get_versions()['version']
del get_versions

def load_ipython_extension(ipython):  # pragma: no cover

    def dis_magic(line, cell=None):
        if cell is None:
            return d(line)
        return d(cell)
    ipython.register_magic_function(dis_magic, 'line_cell', 'dis')

    def ast_magic(line, cell=None):
        if cell is None:
            return a(line)
        return a(cell)
    ipython.register_magic_function(ast_magic, 'line_cell', 'ast')

__all__ = [
    'a',
    'd',
    'display',
    'Code',
    'CodeTransformer',
    'Flag',
    'instructions',
    'matchany',
    'not_',
    'option',
    'or_',
    'pattern',
    'pattern',
    'plus',
    'pformat_ast',
    'pprint_ast',
    'seq',
    'var',
    'transformers',
]