# Jamie Pospishil
# Version 3.6
# 2019-04-08

import pygame
import glob
import os

debug = False

class SFX:
    mixer_channels = 16
    # starting with 16 for mixer_channels, this can be increased if we need to
    number_of_reserved_sound_channels = 2
    # channels 0 - 1 are reserved for possible future use

    # channels_used will hold onto what channels each instance is using
    channels_used = []
    # adds reserved channels to used list
    for i in range(number_of_reserved_sound_channels):
        channels_used.append(i)

    last_channel_used = -1

    def __init__(self, folder_path, file_ext):
        self.sound_play_info = []
        self.sound_file_name = []
        self.sound_channel = []
        self.last_local_used = -1
        self.channels_used = []
        if debug == True:
            self.blep = []
        # self.sound_dictionary = {}
        for ext in range(len(file_ext)):
            for path in range(len(folder_path)):
                self.load_folder(folder_path[path], file_ext[ext])

    def __del__(self):
        for i in self.channels_used:
            SFX.channels_used.remove(i)

    def play(self, sound_passed):
        # if self.sound_file_name.__contains__(soundPassed):
        sound_index = self.sound_file_name.index(sound_passed)
        sound = self.sound_play_info[sound_index]
        if isinstance(sound, pygame.mixer.Sound):
            if debug == True:
                print("Channel", self.blep[sound_index], self.sound_file_name[sound_index],
                      self.sound_play_info[sound_index], self.sound_channel[sound_index])
            self.sound_channel[sound_index].play(sound)
        else:
            pygame.mixer.music.load(sound)
            pygame.mixer.music.play(0)

    # def play_test(self, sound_passed):
    #   sound_to_play = self.sound_dictionary_list[sound_passed]
    #  sound_to_play["channel"].play(sound_to_play["play_info"])

    def setVolume(self, sound_passed, value):
        sound_channel = self.sound_channel[self.sound_file_name.index(sound_passed)]
        if debug:
            print("volume of ", sound_passed, " is ", sound_channel.get_volume())
        sound_channel.set_volume(value)
        if debug:
            print("volume of ", sound_passed, " is ", sound_channel.get_volume())

    def playIfNotBusy(self, sound_passed):
        # if self.sound_file_name.__contains__(soundPassed):
        sound_index = self.sound_file_name.index(sound_passed)
        if not self.sound_channel[sound_index].get_busy():
            sound = self.sound_play_info[sound_index]
            if isinstance(sound, pygame.mixer.Sound):
                if debug == True:
                    print("Channel", self.blep[sound_index], self.sound_file_name[sound_index],
                          self.sound_play_info[sound_index], self.sound_channel[sound_index])
                self.sound_channel[sound_index].play(sound)
            else:
                pygame.mixer.music.load(sound)
                pygame.mixer.music.play(0)

    @staticmethod
    def stop_music():
        """
        Stops playback of the music track. Will not error if there is not a music track.
        :return: None
        """
        pygame.mixer.music.stop()

    @staticmethod
    def stop_sound(sound):
        """
        Stops playback of the specified sound. Will not error if passed an incorrect object.
        If you want to stop ALL sound, call this and stop_music().
        :param sound: The pygame.mixer.Sound object to stop.
        :return: None.
        """
        if isinstance(sound, pygame.mixer.Sound):
            pygame.mixer.Sound.stop(sound)

    def load_folder(self, folder_path, extension):
        """
        Takes all the files from the specified path (folder), does pygame.mixer.Sound(file) on them,
        and then puts them into a list.
        :param folder_path: The relative path of the folder.
        Ex. 'SFX'
        Ex. "SFX/Character"
        :param extension: A file extension. Pygame works best with .wav files.
        By default, this is .wav
        Ex. '.wav'
        :return: A list of Pygame sounds.
        Ex. SFX = loadSoundFolder('SFX', '.wav')
        """
        # https://stackoverflow.com/questions/18262293/how-to-open-every-file-in-a-folder/18262324#18262324

        for filename in glob.glob(os.path.join(folder_path, "*" + extension)):
            path_len = len(folder_path)
            # file_len = len(filename) - path_len - len(extension)
            file_len = len(filename) - path_len
            prefix_len = 0  # remove 001_filename ?
            if extension == ".wav" or extension == ".ogg":
                if "Music" in folder_path:
                    append_str = filename[path_len + 1:(path_len + file_len)]
                    self.sound_file_name.append(append_str)
                    self.sound_play_info.append(filename)
                    self.sound_channel.append(-2)
                    if debug == True:
                        self.blep.append(-2)

                    sound_dictionary = {
                        "type": "Music",
                        "filename": append_str,
                        "play_info": filename
                    }
                    # self.sound_dictionary_list.append(sound_dictionary)
                else:
                    local_channel = int(filename[path_len + 1:path_len + 4], 10)
                    append_str = filename[path_len + 1 + prefix_len:(path_len + file_len)]
                    self.sound_file_name.append(append_str)
                    self.sound_play_info.append(pygame.mixer.Sound(filename))

                    if self.last_local_used == -1:
                        self.last_local_used = local_channel
                        done = False
                        i = SFX.number_of_reserved_sound_channels
                        while not done:
                            if i == SFX.mixer_channels:
                                SFX.mixer_channels *= 2
                                pygame.mixer.set_num_channels(SFX.mixer_channels)
                            if i not in SFX.channels_used:
                                if debug == True:
                                    self.blep.append(i)
                                self.sound_channel.append(pygame.mixer.Channel(i))
                                SFX.channels_used.append(i)
                                self.channels_used.append(i)
                                SFX.last_channel_used = i
                                done = True
                            i += 1

                    elif self.last_local_used == local_channel:
                        self.sound_channel.append(pygame.mixer.Channel(SFX.last_channel_used))
                        if debug == True:
                            self.blep.append(SFX.last_channel_used)

                    else:
                        self.last_local_used = local_channel
                        done = False
                        i = SFX.number_of_reserved_sound_channels
                        while not done:
                            if i == SFX.mixer_channels:
                                SFX.mixer_channels *= 2
                                pygame.mixer.set_num_channels(SFX.mixer_channels)
                            if i not in SFX.channels_used:
                                if debug == True:
                                    self.blep.append(i)
                                self.sound_channel.append(pygame.mixer.Channel(i))
                                SFX.channels_used.append(i)
                                self.channels_used.append(i)
                                SFX.last_channel_used = i
                                done = True
                            i += 1

                    # sound_dictionary[]
            #            sound_dictionary = {
            #               "filename": append_str,
            #              "type": "SFX",
            #             "play_info": filename,
            #            "channel": pygame.mixer.Channel(append_channel)
            #       }
            # self.sound_dictionary_list.append(sound_dictionary)
            elif extension == '.mp3' or extension == ".xm" or extension == ".mod":
                append_str = filename[path_len + 1:(path_len + file_len)]
                self.sound_file_name.append(append_str)
                self.sound_play_info.append(filename)
                self.sound_channel.append(-3)
                self.blep.append(-3)
                sound_dictionary = {
                    "type": "SFX_M",
                    "filename": append_str,
                    "play_info": filename,
                }
                # self.sound_dictionary_list.append(sound_dictionary)

    def load_filenames(self, folder_path, extension):
        """
        Takes all the files from the specified path (folder) and puts their filenames into a list.
        :param folder_path: The file path.
        Ex. "SFX"
        Ex. 'SFX/Character'
        Ex. 'SFX\\Character' # If you want to use backslashes, you have to use 2 because it's an escape character.
        :param extension: A file extension.
        Ex. ".wav"
        :return: A list of strings. The path will precede the filename in brackets.
        Ex. SFX = loadSoundFolder("SFX/Character", ".wav")
        print(SFX[0])   # [SFX] filename.wav
        """
        # https://stackoverflow.com/questions/18262293/how-to-open-every-file-in-a-folder/18262324#18262324
        for filename in glob.glob(os.path.join(folder_path, "*" + extension)):
            if extension == ".wav" or extension == ".ogg":
                if "Music" in folder_path:
                    self.sound_play_info.append(
                        filename.strip(folder_path + "\\").strip(folder_path + "/"))
                else:
                    self.sound_play_info.append(
                        filename.strip(folder_path + "\\").strip(folder_path + "/"))
            elif extension == '.mp3' or extension == ".xm" or extension == ".mod":
                self.sound_play_info.append(filename.strip(folder_path + "\\").strip(folder_path + "/"))

    @staticmethod
    def load_filenames_soundboard(folder_path, extension):
        """
        Takes all the files from the specified path (folder) and puts their filenames into a list.
        :param folder_path: The file path.
        Ex. "SFX"
        Ex. 'SFX/Character'
        Ex. 'SFX\\Character' # If you want to use backslashes, you have to use 2 because it's an escape character.
        :param extension: A file extension.
        Ex. ".wav"
        :return: A list of strings. The path will precede the filename in brackets.
        Ex. SFX = loadSoundFolder("SFX/Character", ".wav")
        print(SFX[0])   # [SFX] filename.wav
        """
        # https://stackoverflow.com/questions/18262293/how-to-open-every-file-in-a-folder/18262324#18262324
        return_list = []
        for filename in glob.glob(os.path.join(folder_path, "*" + extension)):
            if extension == ".wav" or extension == ".ogg":
                if "Music" in folder_path:
                    return_list.append(
                        "[" + folder_path + "] " + filename.strip(folder_path + "\\").strip(folder_path + "/"))
                else:
                    return_list.append(
                        "[" + folder_path + "] " + filename.strip(folder_path + "\\").strip(folder_path + "/"))
            elif extension == '.mp3' or extension == ".xm" or extension == ".mod":
                return_list.append(
                    "[" + folder_path + "] " + filename.strip(folder_path + "\\").strip(folder_path + "/"))
        return return_list

# SoundBoard Code
if __name__ == "__main__":
    # Put Folder Paths here to test sounds
    AllFolderPaths = ['Audio/Music', 'Audio/SFX/Player']
    AllFileExtensions = ['.wav', '.mp3', '.ogg', '.xm', '.mod']

    # pygame.mixer.pre_init(frequency=48000, size=-16, channels=2, buffer=4096)
    pygame.init()
    pygame.mixer.init(frequency=48000, size=-16, channels=2, buffer=4096)
    pygame.mixer.set_num_channels(SFX.mixer_channels)
    pygame.mixer.set_reserved(SFX.number_of_reserved_sound_channels)

    soundboard = SFX(AllFolderPaths, AllFileExtensions)
    myClock = pygame.time.Clock()
    fontConsolas = pygame.font.SysFont("Consolas", 12)

    # LAYOUT STUFF
    screenW = 1280
    screenH = 720
    numCol = 4
    numRow = 8
    # LAYOUT STUFF

    screenRes = (screenW, screenH)
    screen = pygame.display.set_mode(screenRes)
    screenSurf = pygame.Surface(screenRes, pygame.SRCALPHA)

    colBorder = screenH // 50
    rowBorder = colBorder
    totalColBorder = colBorder * (numCol + 1)
    totalRowBorder = rowBorder * (numRow + 1)

    boxWidth = (screenW - totalColBorder) // numCol
    boxHeight = (screenH - totalRowBorder) // numRow

    boxCoords = []
    boxID = []
    for i in range(numCol):
        for j in range(numRow):
            boxCoords.append([i, j])
            boxID.append((j + i * numCol))

    running = True
    mouseClick = [0, 1]
    currentSound = 0
    while running:
        screenSurf.fill((0, 0, 0, 0))
        myClock.tick_busy_loop(60)
        myEvents = pygame.event.get()
        # myKeys = pygame.key.get_pressed()
        mousePress = pygame.mouse.get_pressed()
        mousePos = pygame.mouse.get_pos()

        if not mousePress[0]:
            mouseClick[0] = False
        if not mousePress[1]:
            mouseClick[1] = False

        for event in myEvents:
            if event.type == pygame.QUIT:
                running = False
                break
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    running = False
                    break

        for i in range(len(boxID)):
            outerX = boxCoords[i][0] * (boxWidth + colBorder) + colBorder
            outerY = boxCoords[i][1] * (boxHeight + rowBorder) + rowBorder
            innerX = boxCoords[i][0] * (boxWidth + colBorder) + 2 * colBorder
            innerY = boxCoords[i][1] * (boxHeight + rowBorder) + 2 * rowBorder
            innerWidth = boxWidth - 2 * colBorder
            innerHeight = boxHeight - 2 * rowBorder

            outerColor = [127, 127, 127, 191]
            innerColor = [191, 127, 127, 191]
            if outerX <= mousePos[0] <= (outerX + boxWidth):
                if outerY <= mousePos[1] <= (outerY + boxHeight):
                    outerColor = [255, 255, 191, 191]
                    if mousePress[0]:
                        innerColor = [127, 127, 191, 255]
                        if i < len(soundboard.sound_play_info) and not mouseClick[0]:
                            # soundboard.stop_music()
                            # soundboard.stop_sound(soundboard.sound_play_info[currentSound])
                            soundboard.play(soundboard.sound_file_name[i])
                            currentSound = i
                            mouseClick[0] = True

            pygame.draw.rect(screenSurf, outerColor, (outerX, outerY, boxWidth, boxHeight))
            pygame.draw.rect(screenSurf, innerColor, (innerX, innerY, innerWidth, innerHeight))
            if i < len(soundboard.sound_file_name):
                text = fontConsolas.render("{}".format("[" + (str(i)) + "]" + soundboard.sound_file_name[i]), False,
                                           (255, 255, 255))
                screenSurf.blit(text, (innerX + colBorder // 2, innerY + (innerHeight // 2 - rowBorder // 2)))

        screen.blit(screenSurf, (0, 0, screenW, screenH))
        pygame.display.update()

    pygame.quit()


#!/usr/bin/env python3

'''
*****************************************************************************************
*
*        		===============================================
*           		    HolA Bot (HB) Theme (eYRC 2022-23)
*        		===============================================
*
*  This script should be used to implement Task 0 of HolA Bot (KB) Theme (eYRC 2022-23).
*
*  This software is made available on an "AS IS WHERE IS BASIS".
*  Licensee/end user indemnifies and will keep e-Yantra indemnified from
*  any and all claim(s) that emanate from the use of the Software or
*  breach of the terms of this agreement.
*
*****************************************************************************************
'''

# Team ID:			2339
# Author List:		Dhruva Shaw
# Filename:			task_0.py
# Functions:
# 					main, update_pose, equationroots, make_semicircle, give_pos_y
# Nodes:		    turtlebot_controller

####################### IMPORT MODULES #######################
import sys
import traceback
import rospy
from geometry_msgs.msg import Twist
from turtlesim.msg import Pose
from math import sqrt, pi
from typing import Union, Tuple
##############################################################

class TurtleBot:

    def __init__(self) -> None:
        """
        Initializes the important variables required in the whole class
        """

        # Creates a node with name 'turtlebot_controller' and make sure it is a
        # unique node (using anonymous=True).
        rospy.init_node('turtlebot_controller', anonymous=True)

        # Publisher which will publish to the topic '/turtle1/cmd_vel'.
        self.velocity_publisher = rospy.Publisher('/turtle1/cmd_vel',Twist, queue_size=10)

        # A subscriber to the topic '/turtle1/pose'. self.update_pose is called
        # when a message of type Pose is received.
        self.pose_subscriber = rospy.Subscriber('/turtle1/pose',Pose, self.update_pose)

        #Storing the position of the turtle
        self.pose = Pose()

        #The frequency of the turle
        self.rate = rospy.Rate(10)
        self.vel = Twist()

        #Final position of y required to make a semicircle
        self.final_pos_circle_y = 0

        #Initial startin point of the turtle
        self.start_point_y = self.pose.y
        self.start_point_x = self.pose.x

    def update_pose(self, data: Pose) -> None:
        """
        Callback function which is called when a new message of type Pose is
        received by the subscriber.
        """
        #saving the individual data in the Pose class
        self.pose = data

        #Storing the rounded data up to two decimal place as it is easier to calculate
        self.pose.x = round(self.pose.x,2)
        self.pose.y = round(self.pose.y,2)

    def equationroots(self, a: int, b: int, c: int) -> Tuple[Union[int, float]]:
        """Takes a,b and c coefficients of the root and calculates the roots using shreedharacharya's formula

        Returns:
            Tuple[Union[int, float]]: The Root of the equation
        """
        # calculating discriminant using formula
        dis = b * b - 4 * a * c
        #returning both the values using shreedharacharya's formula
        return (-b + sqrt(dis))/(2*a) , (-b - sqrt(dis))/(2*a)

    def give_pos_y(self, y_pos: int, radius: int) -> Union[int, float]:
        """Gives the required position of the coordinate that the turtle needs to move to form a semicircle from initial y axis and radius

        Args:
            y_pos (int): The initial y position of the tutrle at the start
            radius (int): THe radius of the semicircle

        Returns:
            Union[int, float]: Required position of the coordinate that the turtle needs to move to form a semicircle
        """
        y1, y2 = self.equationroots(1, -2*y_pos, -(((radius*2)**2)-(y_pos**2)))
        return y1 if y1 > 0 else y2

    def make_semicircle(self, radius: int) -> None:
        """Make a D shape in teh turtlesim

        Args:
            radius (int): The radius of 'D' or semicircle
        """

        #semicircle
        while not rospy.is_shutdown():
            #Initialization of the velocity variables
            self.vel.linear.x: int = radius #Since x coordinate needs to be moved to make a D shape
            self.vel.linear.y: int = 0
            self.vel.linear.z: int = 0
            self.vel.angular.x: int = 0
            self.vel.angular.y: int = 0
            self.vel.angular.z: int = radius #also to make a curve the angulare z needs to move

            #Logging
            rospy.loginfo("Radius = %f",radius)
            rospy.loginfo("x = %f",self.pose.x)
            rospy.loginfo("y = %f",self.pose.y)
            rospy.loginfo("Final position for the semicircle without the bottom horizontal line = %f",self.final_pos_circle_y)

            #Getting some impt values just at the startup
            if self.pose.x != 0 and self.pose.y != 0 and self.final_pos_circle_y == 0:
                self.final_pos_circle_y = self.give_pos_y(self.pose.y, radius) #final position of y to make semicircle
                self.start_point_y = self.pose.y #inital position of y
                self.start_point_x = self.pose.x #inital position of x

            #While loop breaking logic to make the semicircle
            if self.pose.y >= self.final_pos_circle_y and self.final_pos_circle_y != 0:
                self.vel.linear.x: int = 0
                self.vel.angular.z = 0
                self.velocity_publisher.publish(self.vel) #Publishes the velocity to the turtle
                break
            self.velocity_publisher.publish(self.vel)
            self.rate.sleep() #This function properply controles the frequency of the turtle

        #Rotation
        self.vel.linear.x = 0 #Seeting this to 0 as it is not required
        self.vel.angular.z=abs(90*2*pi/360) #angles to radians
        self.velocity_publisher.publish(self.vel)
        self.rate.sleep()

        #Linear motion to complete the D
        while not rospy.is_shutdown():
            self.vel.linear.x = 0
            self.vel.linear.y = sqrt((self.pose.y - self.start_point_y)**2 + (self.pose.x - self.start_point_x)**2) #Using the distance formula to calculate the exact linear distance that needs to be cover
            self.vel.linear.z = 0
            self.vel.angular.x = 0
            self.vel.angular.y = 0
            self.vel.angular.z = 0

            #Breaking logic of the linear motion
            if self.pose.y <= self.start_point_y:
                self.vel.linear.y = 0
                self.velocity_publisher.publish(self.vel)
                break

            rospy.loginfo("x = %f",self.pose.x)
            rospy.loginfo("y = %f",self.pose.y)

            self.velocity_publisher.publish(self.vel)

            self.rate.sleep()

        rospy.spin()

def main() -> None:
    """The main function which runs the whole program
    """
    #Init of the class
    turtlebot = TurtleBot()

    #Making the D shape of radius 1
    turtlebot.make_semicircle(1)

######### YOU ARE NOT ALLOWED TO MAKE CHANGES TO THIS PART #########
if __name__ == "__main__":
    try:
        print("------------------------------------------")
        print("         Python Script Started!!          ")
        print("------------------------------------------")
        main()

    except:
        print("------------------------------------------")
        traceback.print_exc(file=sys.stdout)
        print("------------------------------------------")
        sys.exit()

    finally:
        print("------------------------------------------")
        print("    Python Script Executed Successfully   ")
        print("------------------------------------------")


import argparse
from datetime import datetime

import torch
import torch.nn.functional as F
from torch.nn.utils import clip_grad_norm_
from torchmetrics.image.fid import FrechetInceptionDistance
import numpy as np
import os

import time

from torchvision.transforms import (
    CenterCrop,
    Compose,
    InterpolationMode,
    RandomHorizontalFlip,
    Resize,
    ToTensor,
)
from datasets import load_dataset
from diffusers.optimization import get_scheduler
from tqdm.auto import tqdm
from torchinfo import summary

from scheduler import DDIMScheduler
from model import UNet
from utils import save_images, normalize_to_neg_one_to_one, plot_losses
from datetime import datetime
import bisect

n_timesteps = 1000
n_inference_timesteps = 50

def main(args):
    # Setup Code
    timestamp = datetime.now().strftime('%Y-%m-%d-%H-%M-%S')

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    augmentations = Compose([
        Resize(args.resolution, interpolation=InterpolationMode.BILINEAR),
        CenterCrop(args.resolution),
        RandomHorizontalFlip(),
        ToTensor(),
    ])

    if args.dataset_name is not None:
        dataset = load_dataset(
            args.dataset_name,
            args.dataset_config_name,
            cache_dir=args.cache_dir,
            split="train",
        )
    else:
        dataset = load_dataset("imagefolder",
                                data_dir=args.train_data_dir,
                                cache_dir=args.cache_dir,
                                split="train")

    def transforms(examples):
        images = [augmentations(image.convert("RGB")) for image in examples["image"]]
        return {"input": images}

    dataset.set_transform(transforms)

    train_dataloader = torch.utils.data.DataLoader(dataset,
                                                   batch_size=args.train_batch_size,
                                                   shuffle=True)

    fid_dataloader = torch.utils.data.DataLoader(dataset,
                                                   batch_size=4096,
                                                   shuffle=False,
                                                    )
    real_image_samples = next(iter(train_dataloader))['input']

    # LTH Code
    prune_percent_per_LTH_round = 1 - (1 - args.LTH_prune_percent)**(1/args.LTH_rounds)
    print("Pruning", str(round(prune_percent_per_LTH_round, 4)*100) + "% for a total prune percent of", str(args.LTH_prune_percent*100) + "% over", args.LTH_rounds, "rounds")

    model = UNet(3, image_size=args.resolution, hidden_dims=[16, 32, 64, 128])
    model = model.to(device)

    mask = dict()
    init_weights = dict()

    for name, param in model.named_parameters():
        mask[name] = torch.ones(param.shape).to(device)
        init_weights[name] = param.detach().clone().to(device)

    for rnd in range(args.LTH_rounds):
        print("LTH round:", rnd)

        # Copy masked initial weights to model
        with torch.no_grad():
            for name, param in model.named_parameters():
                param.copy_(mask[name]*init_weights[name])

        model.train()
        noise_scheduler = DDIMScheduler(num_train_timesteps=n_timesteps,
                                        beta_schedule="cosine", device=device)

        optimizer = torch.optim.AdamW(
            model.parameters(),
            lr=args.learning_rate,
            betas=(args.adam_beta1, args.adam_beta2),
            weight_decay=args.adam_weight_decay,
            eps=args.adam_epsilon,
        )

        lr_scheduler = get_scheduler(
            args.lr_scheduler,
            optimizer=optimizer,
            num_warmup_steps=args.lr_warmup_steps,
            num_training_steps=(len(train_dataloader) * args.num_epochs) //
            args.gradient_accumulation_steps,
        )

        for epoch in range(args.LTH_epochs_per_rounds):
            progress_bar = tqdm(total=len(train_dataloader))
            progress_bar.set_description(f"Epoch {epoch}")
            for step, batch in enumerate(train_dataloader):
                optimizer.zero_grad()
                clean_images = batch["input"].to(device)
                clean_images = normalize_to_neg_one_to_one(clean_images)

                batch_size = clean_images.shape[0]
                noise = torch.randn(clean_images.shape).to(device)
                timesteps = torch.randint(0, noise_scheduler.num_train_timesteps,
                                         (batch_size,), device=device).long()
                noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps)

                noise_pred = model(noisy_images, timesteps)["sample"]
                loss = F.l1_loss(noise_pred, noise)
                loss.backward()

                if args.use_clip_grad:
                    clip_grad_norm_(model.parameters(), 1.0)

                # Apply our 0/1 mask
                for name, param in model.named_parameters():
                    param.grad *= mask[name]

                optimizer.step()
                lr_scheduler.step()

                progress_bar.update(1)
            progress_bar.close()
        # Identify threshold and update mask
        all_weights = list()

        for name, param in model.named_parameters():
            filtered_params = torch.abs(param*mask[name])
            all_weights.extend(filtered_params.flatten().tolist())

        sorted_arr = sorted(all_weights)

        #print(len(sorted_arr), bisect.bisect(sorted_arr, 0))
        sorted_arr = sorted_arr[bisect.bisect(sorted_arr, 0):]

        #print(len(sorted_arr))

        threshold = sorted_arr[int(len(sorted_arr)*prune_percent_per_LTH_round)]

        for name, param in model.named_parameters():
            mask[name] = mask[name]*(torch.abs(param) > threshold).int()

    # Actual Model training code
    model = UNet(3, image_size=args.resolution, hidden_dims=[16, 32, 64, 128])

    # Copy masked initial weights to model
    with torch.no_grad():
        for name, param in model.named_parameters():
            param.copy_(mask[name]*init_weights[name])

    noise_scheduler = DDIMScheduler(num_train_timesteps=n_timesteps,
                                    beta_schedule="cosine", device=device)

    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=args.learning_rate,
        betas=(args.adam_beta1, args.adam_beta2),
        weight_decay=args.adam_weight_decay,
        eps=args.adam_epsilon,
    )

    lr_scheduler = get_scheduler(
        args.lr_scheduler,
        optimizer=optimizer,
        num_warmup_steps=args.lr_warmup_steps,
        num_training_steps=(len(train_dataloader) * args.num_epochs) //
        args.gradient_accumulation_steps,
    )

    if not os.path.exists("logs"):
        os.mkdir("logs")

    log_filename = "logs/" +str(args.LTH_prune_percent) + "_log_" + str(time.time_ns() // 1_000_000) + ".csv"
    with open(log_filename, "a") as file:
        file.write("epoch,fid,train_loss\n")

    model = model.to(device)
    summary(model, [(1, 3, args.resolution, args.resolution), (1,)], verbose=1)

    global_step = 0
    losses = []
    for epoch in range(args.num_epochs):
        model.train()
        progress_bar = tqdm(total=len(train_dataloader))
        progress_bar.set_description(f"Epoch {epoch}")
        losses_log = 0
        for step, batch in enumerate(train_dataloader):
            clean_images = batch["input"].to(device)
            clean_images = normalize_to_neg_one_to_one(clean_images)

            batch_size = clean_images.shape[0]
            noise = torch.randn(clean_images.shape).to(device)
            timesteps = torch.randint(0, noise_scheduler.num_train_timesteps,
                                     (batch_size,), device=device).long()
            noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps)

            noise_pred = model(noisy_images, timesteps)["sample"]
            loss = F.l1_loss(noise_pred, noise)
            loss.backward()

            if args.use_clip_grad:
                clip_grad_norm_(model.parameters(), 1.0)

            # Apply our 0/1 mask
            for name, param in model.named_parameters():
                param.grad *= mask[name]

            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()

            progress_bar.update(1)
            losses_log += loss.detach().item()
            logs = {
                "loss_avg": losses_log / (step + 1),
                "loss": loss.detach().item(),
                "lr": lr_scheduler.get_last_lr()[0],
                "step": global_step
            }

            progress_bar.set_postfix(**logs)
            global_step += 1
        progress_bar.close()
        losses.append(losses_log / (step + 1))

        # Generate sample images for visual inspection
        if epoch % args.save_model_epochs == 0:
            with torch.no_grad():
                # has to be instantiated every time, because of reproducibility
                generator = torch.manual_seed(0)
                generated_images = noise_scheduler.generate(
                    model,
                    num_inference_steps=n_inference_timesteps,
                    generator=generator,
                    eta=1.0,
                    use_clipped_model_output=True,
                    batch_size=args.eval_batch_size,
                    output_type="numpy")

                fid = FrechetInceptionDistance(feature=64, normalize=False)

                fid.update(torch.tensor(256*np.moveaxis(generated_images['sample'], -1, 1)).to(torch.uint8), real=False)
                fid.update((256*real_image_samples).to(torch.uint8), real=True)

                with open(log_filename, "a") as file:
                    file.write(str(epoch) + "," + str(fid.compute().item()) + "," + str(losses[-1]) + "\n")

                save_images(generated_images, epoch, args, title=str(args.LTH_prune_percent) + "_" + str(epoch))
                plot_losses(losses, f"{args.loss_logs_dir}_{timestamp}/{epoch}/")

                #if not os.path.exists("trained_models"):
                #    os.mkdir("trained_models")

                #torch.save(
                #    {
                #        'model_state': model.state_dict(),
                #        'optimizer_state': optimizer.state_dict(),
                #    }, args.output_dir)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Simple example of a training script.")
    parser.add_argument("--dataset_name", type=str, default="flowers")
    parser.add_argument("--dataset_config_name", type=str, default=None)
    parser.add_argument("--train_data_dir",
                        type=str,
                        default=None,
                        help="A folder containing the training data.")
    parser.add_argument("--output_dir", type=str, default="trained_models/ddpm-model-64.pth")
    parser.add_argument("--samples_dir", type=str, default="replica_samples/")
    parser.add_argument("--loss_logs_dir", type=str, default="replica_training_logs")
    parser.add_argument("--cache_dir", type=str, default=None)
    parser.add_argument("--resolution", type=int, default=64)
    parser.add_argument("--train_batch_size", type=int, default=64)
    parser.add_argument("--eval_batch_size", type=int, default=32)
    parser.add_argument("--num_epochs", type=int, default=2000)
    parser.add_argument("--save_model_epochs", type=int, default=10)
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1)
    parser.add_argument("--learning_rate", type=float, default=1e-3)
    parser.add_argument("--lr_scheduler", type=str, default="cosine")
    parser.add_argument("--lr_warmup_steps", type=int, default=500)
    parser.add_argument("--adam_beta1", type=float, default=0.9)
    parser.add_argument("--adam_beta2", type=float, default=0.999)
    parser.add_argument("--adam_weight_decay", type=float, default=0)
    parser.add_argument("--adam_epsilon", type=float, default=1e-7)
    parser.add_argument("--use_clip_grad", type=bool, default=False)
    parser.add_argument("--logging_dir", type=str, default="logs")

    parser.add_argument("--LTH_rounds", type=int, default=5)
    parser.add_argument("--LTH_epochs_per_rounds", type=int, default=10)
    parser.add_argument("--LTH_prune_percent", type=float, default=0.2)

    args = parser.parse_args()

    if args.dataset_name is None and args.train_data_dir is None:
        raise ValueError(
            "You must specify either a dataset name from the hub or a train data directory.")

    main(args)


from __future__ import print_function
import httplib2
import os

from apiclient import discovery
import oauth2client
from oauth2client import client
from oauth2client import tools

import json

courses_filename = "courses.json"
courseCalendarID = "ucdavis.edu_t1v777vkdgdrrd1s1doltuvtd4@group.calendar.google.com"

try:
    import argparse
    flags = argparse.ArgumentParser(parents=[tools.argparser]).parse_args()
except ImportError:
    flags = None

# If modifying these scopes, delete your previously saved credentials
# at ~/.credentials/calendar-python-quickstart.json
SCOPES = 'https://www.googleapis.com/auth/calendar'
CLIENT_SECRET_FILE = 'client_secret.json'
APPLICATION_NAME = 'Google Calendar API Python Quickstart'

# Got this code from sample, haven't really looked into it yet.
def get_credentials():
    """Gets valid user credentials from storage.

    If nothing has been stored, or if the stored credentials are invalid,
    the OAuth2 flow is completed to obtain the new credentials.

    Returns:
        Credentials, the obtained credential.
    """
    home_dir = os.path.expanduser('~')
    credential_dir = os.path.join(home_dir, '.credentials')
    if not os.path.exists(credential_dir):
        os.makedirs(credential_dir)
    credential_path = os.path.join(credential_dir,
                                   'calendar-python-quickstart.json')

    store = oauth2client.file.Storage(credential_path)
    credentials = store.get()
    if not credentials or credentials.invalid:
        flow = client.flow_from_clientsecrets(CLIENT_SECRET_FILE, SCOPES)
        flow.user_agent = APPLICATION_NAME
        if flags:
            credentials = tools.run_flow(flow, store, flags)
        else: # Needed only for compatibility with Python 2.6
            credentials = tools.run(flow, store)
        print('Storing credentials to ' + credential_path)
    return credentials

# Clears calendar and adds courses. Recommend doing this with empty calendar,
# then checking results before importing
def main():
    """Shows basic usage of the Google Calendar API.

    Creates a Google Calendar API service object and outputs a list of the next
    10 events on the user's calendar.
    """
    credentials = get_credentials()
    http = credentials.authorize(httplib2.Http())
    service = discovery.build('calendar', 'v3', http=http)

    print('Clearing calendar')
    eventsResult = service.events().list(
        calendarId=courseCalendarID, timeMin='2016-01-01T00:00:00-07:00', singleEvents=True,
        orderBy='startTime').execute()
    events = eventsResult.get('items', [])
    del_request = service.new_batch_http_request()

    for event in events:
        del_request.add(service.events().delete(calendarId=courseCalendarID, eventId=event["id"]))

    del_request.execute()

    with open(courses_filename) as courses_file:
        add_request = service.new_batch_http_request()
        courses = json.load(courses_file)
        for course in courses.values():
            add_request.add(service.events().insert(calendarId=courseCalendarID, body=course))

    add_request.execute()

if __name__ == '__main__':
    main()


import base64
import hashlib
import hmac
import requests
import time
import json

from rest_framework.exceptions import ValidationError

from Troy.settings import base
from apps.oauth.models import Auth

class GoogleService(object):
    def __init__(self):
        self.google_id_token_info_url = 'https://www.googleapis.com/oauth2/v3/tokeninfo'
        self.google_email_info_url = 'https://www.googleapis.com/oauth2/v1/tokeninfo'

    def google_validate_email(self, data):
        # [access_token]으로 Google OAuth에 request 전송
        # email로 user 유효성 검증 (email은 unique한 값을 가진다는 성질을 이용한 logic)
        context = {
            'access_token': data['access_token']
        }
        response = requests.post(self.google_email_info_url, data=context)
        if not response.ok:
            raise ValidationError(response.reason)

        response_json = response.json()
        user_email = data['email']

        response_flag = [user_email == response_json['email']]
        return response_flag

    def google_validate_client_id(self, data):
        # [id_token]으로 Google OAuth에 request 전송
        # client_id값으로 user 유효성 검증 (request를 보낸 client가 우리의 client가 맞는지 확인하는 logic)
        context = {
            'id_token': data['id_token']
        }
        response = requests.post(self.google_id_token_info_url, data=context)
        if not response.ok:
            raise ValidationError(response.reason)

        response_json = response.json()
        google = getattr(base, 'auth')['google']
        client_id = google['client_id']

        response_flag = [client_id == response_json['aud']]
        return response_flag

class AuthService(object):
    def __init__(self, data):
        self.oauth = data

    @staticmethod
    def authenticate_provider(provider):
        # social login 확장 고려하여 provider 구분
        if provider == 'google':
            return True
        else:
            return False

    def sef_auth_info(self, **kwargs):
        auth = Auth(
            oauth_type=self.oauth['oauth_type'],
            oauth_token=self.oauth['oauth_token']
        )
        return auth

class AuthSMSService(object):
    def __init__(self):
        self.sender = '01089099796'
        self.naver = getattr(base, 'auth')['naver']
        self.url = 'https://sens.apigw.ntruss.com'
        self.uri = f'/sms/v2/services/{self.naver["service_id"]}/messages'

    def make_signature(self, timestamp):
        secret_key = bytes(self.naver['secret_key'], 'UTF-8')

        message = "POST "+self.uri+"\n"+timestamp+"\n"+self.naver['access_key']
        message = bytes(message, 'UTF-8')
        signingKey = base64.b64encode(hmac.new(secret_key, message, digestmod=hashlib.sha256).digest())

        return signingKey

    def send_sms(self, phone_number, auth_number):
        timestamp = str(int(time.time()*1000))
        context = {
            'type': 'SMS',
            'contentType': 'COMM',
            'countryCode': '82',
            'from': self.sender,
            'content': f'Troy 인증번호 [{auth_number}]를 입력해주세요.',
            'messages': [{
                'to': phone_number
            }]
        }
        headers = {
            'Content-Type': 'application/json; charset=utf-8',
            'x-ncp-apigw-timestamp': timestamp,
            'x-ncp-iam-access-key': self.naver['access_key'],
            'x-ncp-apigw-signature-v2': self.make_signature(timestamp=timestamp)
        }
        response = requests.post(self.url+self.uri, headers=headers, json=context)

        if not response.ok:
            raise ValidationError(response.reason)

        return True


import requests, json
from bs4 import BeautifulSoup
from datetime import datetime

def get_cs_notice(row, logger):
    _id, last_num = row['id'], row['last_num']
    result = []
    URL_list = ["http://cs.skku.edu/ajax/board/view/notice/", "http://cs.skku.edu/ajax/board/view/news/",
            "http://cs.skku.edu/ajax/board/view/seminar/", "http://cs.skku.edu/ajax/board/view/recruit/" ]
    category_list = ["notice", "news", "seminar", "recruit"]
    count = 0
    last_num = last_num + 1
    while True:
        URL = URL_list[count] + str(last_num)
        response = requests.post(URL)
        response.encoding = 'UTF-8'
        if response.status_code != requests.codes.ok:
            count = count + 1
            if count == 4:
                break
            continue

        jsonify_response = json.loads(response.text)

        if jsonify_response['result']:
            notice_json = json.loads(response.text)["post"]

            soup = BeautifulSoup(notice_json['text'], 'html5lib')
            img_src = ""
            for img in soup.findAll('img'):
                img_src = img_src + "http://cs.skku.edu" + img['src'] + "#"
            img_src = img_src[:len(img_src)-1]
            contents = soup.text
            link = "http://cs.skku.edu/open/" + category_list[count] + "/view/" + str(notice_json['id'])
            notice = {
                'id': _id,
                'title': notice_json['title'],
                'last_num': notice_json['id'],
                'contents': contents,
                'link': link,
                'img_src': img_src,
                'type': row['type'],
                'ntype': 'H'
            }
            result.append(notice)
            count = 0
            last_num = last_num + 1
        else:
            break
    logger.info('get '+ str(len(result)) + ' new cs notice')
    return result


from jax._src.prng import _bit_stats
from jax._src.tree_util import tree_structure
from jax.experimental.stax import Dense
import neural_tangents as nt
from neural_tangents import stax
import matplotlib.pyplot as plt
import jax.numpy as np
from seaborn.utils import saturate

from tqdm import tqdm

from jax import random
from jax.experimental import optimizers
from jax import jit, grad, vmap

import functools

from IPython.display import set_matplotlib_formats
import matplotlib
from neural_tangents.utils.utils import named_tuple_factory
import seaborn as sns
sns.set_style("darkgrid", {"axes.facecolor": ".95"})
sns.set()

def loss_fn(predict_fn, ys, t, xs=None):
    mean, cov = predict_fn(t=t, get='ntk', x_test=xs, compute_cov=True)
    mean = np.reshape(mean, mean.shape[:1] + (-1,))
    var = np.diagonal(cov, axis1=1, axis2=2)
    ys = np.reshape(ys, (1, -1))

    mean_predictions = 0.5 * np.mean(ys ** 2 - 2 * mean * ys + var + mean ** 2,
                                     axis=1)

    return mean_predictions

# Define train datas
train_points = 5
test_points = 50
noise_scale = 1e-1

def target_fn(x): return np.sin(x)

key = random.PRNGKey(397)
key, x_key, y_key = random.split(key=key, num=3)

# generate test data

train_xs = random.uniform(x_key, (train_points, 1),
                          minval=-np.pi, maxval=np.pi)
train_ys = target_fn(train_xs)
train_ys += noise_scale * random.normal(
    y_key, (train_points, 1)
)

train = (train_xs, train_ys)

test_xs = np.linspace(
    -np.pi, np.pi, test_points
)
test_xs = np.reshape(
    test_xs, (test_points, 1)
)

test_ys = target_fn(test_xs)
test = (test_xs, test_ys)

def plot_fn(train, test, *fs):
    """
    データプロット用関数
    """
    train_xs, train_ys = train
    plt.plot(
        train_xs, train_ys, "ro", markersize=10, label="train"
    )
    if test != None:
        test_xs, test_ys = test
        plt.plot(test_xs, test_ys, 'k--', linewidth=3, label='$f(x)$')

    for f in fs:
        plt.plot(test_xs, f(test_xs), '-', linewidth=3)

    plt.xlim([-np.pi, np.pi])
    plt.ylim([-1.5, 1.5])
    plt.xlabel("x")
    plt.ylabel("f(x)")

plt.figure(figsize=(6, 4))
plot_fn(train, test)
plt.legend(loc="upper left")
plt.savefig("./plot/TrainDatas")

init_fn, apply_fn, kernel_fun = stax.serial(
    stax.Dense(512, W_std=1.5, b_std=0.05), stax.Erf(),
    stax.Dense(512, W_std=1.5, b_std=0.05), stax.Erf(),
    stax.Dense(1, W_std=1.5, b_std=0.05)
)

apply_fn = jit(apply_fn)
kernel_fun = jit(kernel_fun, static_argnums=(2,))

prior_draws = []

for _ in range(10):
    key, net_key = random.split(key)
    _, params = init_fn(net_key, (-1, 1))
    prior_draws += [apply_fn(
        params, test_xs
    )]

plot_fn(train, test)

for p in prior_draws:
    plt.plot(test_xs, p, linewidth=3, alpha=0.7)
    plt.legend(['train', '$f(x)$', 'random draw'], loc='upper left')

plt.savefig("./plot/firstOutput")

kernel = kernel_fun(test_xs, test_xs, "nngp")
std_dev = np.sqrt(np.diag(kernel))

plot_fn(train, test)

plt.fill_between(np.reshape(test_xs, (-1,)), 2 *
                 std_dev, -2 * std_dev, alpha=0.4)

for p in prior_draws:
    plt.plot(test_xs, p, linewidth=3, alpha=0.5)
plt.savefig("./plot/nngpOutput")

# NTK
perdict_fun = nt.predict.gradient_descent_mse_ensemble(
    kernel_fun, train_xs, train_ys, diag_reg=1e-4
)

ntk_mean, ntk_covariance = perdict_fun(
    x_test=test_xs, get='ntk', compute_cov=True
)
print(ntk_mean)
ntk_mean = np.reshape(ntk_mean, (-1,))
print(ntk_mean)
ntk_std = np.sqrt(np.diag(ntk_covariance))

plot_fn(train, test)

plt.plot(test_xs, ntk_mean, 'b-', linewidth=3)
plt.fill_between(
    np.reshape(test_xs, (-1)),
    ntk_mean - 2 * ntk_std,
    ntk_mean + 2 * ntk_std,
    color='blue', alpha=0.2)

plt.xlim([-np.pi, np.pi])
plt.ylim([-1.5, 1.5])

plt.legend(['Train', 'f(x)', 'Bayesian Inference', 'Gradient Descent'],
           loc='upper left')

plt.savefig("./plot/ntk_prediction")

# inference finite time
ts = np.arange(
    0, 10 ** 3, 10**-1
)

print(ts)

ntk_train_loss_mean = loss_fn(perdict_fun, train_ys, ts)
ntk_test_loss_mean = loss_fn(perdict_fun, test_ys, ts, test_xs)

plt.subplot(1, 2, 1)

plt.loglog(ts, ntk_train_loss_mean, linewidth=3)
plt.loglog(ts, ntk_test_loss_mean, linewidth=3)

plt.legend(['Infinite Train', 'Infinite Test'])

plt.subplot(1, 2, 2)

plot_fn(train, None)

plt.plot(test_xs, ntk_mean, 'b-', linewidth=3)
plt.fill_between(
    np.reshape(test_xs, (-1)),
    ntk_mean - 2 * ntk_std,
    ntk_mean + 2 * ntk_std,
    color='blue', alpha=0.2)

plt.legend(
    ['Train', 'Infinite Network'],
    loc='upper left')

plt.savefig("./plot/train_log")

# comopare to train finite network
learning_rate = 0.1
training_steps = 10000

opt_init, opt_update, get_params = optimizers.sgd(
    learning_rate
)
opt_update = jit(opt_update)
loss = jit(lambda params, x, y: 0.5 * np.mean((apply_fn(params, x) - y)**2))
grad_loss = jit(lambda state, x, y: grad(loss)(get_params(state), x, y))

train_losses = []
test_losses = []

opt_state = opt_init(params)

for i in range(training_steps):
    opt_state = opt_update(i, grad_loss(
        opt_state, *train
    ), opt_state)

    train_losses += [loss(get_params(opt_state), *train)]
    test_losses += [loss(get_params(opt_state), *test)]

plt.subplot(1, 2, 1)

plt.loglog(ts, ntk_train_loss_mean, linewidth=3)
plt.loglog(ts, ntk_test_loss_mean, linewidth=3)

plt.loglog(ts, train_losses, 'k-', linewidth=2)
plt.loglog(ts, test_losses, 'k-', linewidth=2)

plt.legend(['Infinite Train', 'Infinite Test', 'Finite'])

plt.subplot(1, 2, 2)

plot_fn(train, None)

plt.plot(test_xs, ntk_mean, 'b-', linewidth=3)
plt.fill_between(
    np.reshape(test_xs, (-1)),
    ntk_mean - 2 * ntk_std,
    ntk_mean + 2 * ntk_std,
    color='blue', alpha=0.2)

plt.plot(test_xs, apply_fn(get_params(opt_state), test_xs), 'k-', linewidth=2)

plt.legend(
    ['Train', 'Infinite Network', 'Finite Network'],
    loc='upper left')

plt.savefig("./plot/compare_finite_network")

# emsanble train
def train_network(key):
    train_losses = []
    test_losses = []

    _, params = init_fn(key, (-1, 1))
    opt_state = opt_init(params)

    for i in tqdm(range(training_steps)):
        train_losses += [np.reshape(
            loss(get_params(opt_state), *train), (1,)
        )]

        test_losses += [np.reshape(
            loss(get_params(opt_state), *test), (1,)
        )]

        opt_state = opt_update(i, grad_loss(
            opt_state, *train
        ), opt_state)

    train_losses = np.concatenate(train_losses)
    test_losses = np.concatenate(test_losses)

    return get_params(
        opt_state
    ), train_losses, test_losses

# params, train_loss, test_loss = train_network(
#     key
# )

# ensamble_size = 10
# ensemble_key = random.split(
#     key, ensamble_size
# )

# params, train_loss, test_loss = vmap(
#     train_network
# )(ensemble_key)

# plt.subplot(1, 2, 1)

# mean_train_loss = np.mean(train_loss, axis=0)
# mean_test_loss = np.mean(test_loss, axis=0)

# plt.loglog(ts, ntk_train_loss_mean, linewidth=3)
# plt.loglog(ts, ntk_test_loss_mean, linewidth=3)

# plt.loglog(ts, mean_train_loss, 'k-', linewidth=2)
# plt.loglog(ts, mean_test_loss, 'k-', linewidth=2)

# plt.xlim([10 ** 0, 10 ** 3])

# plt.xscale('log')
# plt.yscale('log')

# plt.legend(['Infinite Train', 'Infinite Test', 'Finite Ensemble'])

# plt.subplot(1, 2, 2)

# plot_fn(train, None)

# plt.plot(test_xs, ntk_mean, 'b-', linewidth=3)
# plt.fill_between(
#     np.reshape(test_xs, (-1)),
#     ntk_mean - 2 * ntk_std,
#     ntk_mean + 2 * ntk_std,
#     color='blue', alpha=0.2)

# ensemble_fx = vmap(apply_fn, (0, None))(params, test_xs)

# mean_fx = np.reshape(np.mean(ensemble_fx, axis=0), (-1,))
# std_fx = np.reshape(np.std(ensemble_fx, axis=0), (-1,))

# plt.plot(test_xs, mean_fx - 2 * std_fx, 'k--', label='_nolegend_')
# plt.plot(test_xs, mean_fx + 2 * std_fx, 'k--', label='_nolegend_')
# plt.plot(test_xs, mean_fx, linewidth=2, color='black')

# plt.legend(['Train', 'Infinite Network', 'Finite Ensemble'], loc='upper left')

# plt.xlim([-np.pi, np.pi])
# plt.ylim([-1.5, 1.5])

# plt.savefig('./plot/compare_to_ensamble')

# Define other architectures
ResBlock = stax.serial(
    stax.FanOut(2),
    stax.parallel(
        stax.serial(
            stax.Erf(),
            stax.Dense(512, W_std=1.1, b_std=0)
        ),
        stax.Identity()
    ),
    stax.FanInSum()
)

init_fn, apply_fn, kernel_fn = stax.serial(
    stax.Dense(512, W_std=1, b_std=0),
    ResBlock, ResBlock, stax.Erf(),
    stax.Dense(1, W_std=1.5, b_std=0)
)

apply_fn = jit(apply_fn)
kernel_fn = jit(kernel_fn, static_argnums=(2,))

ensamble_size = 5
learning_rate = 0.1
ts = np.arange(
    0, 10**3, learning_rate
)

opt_init, opt_update, get_params = optimizers.sgd(
    learning_rate
)

opt_update = jit(opt_update)

key, = random.split(key, 1)
ensemble_key = random.split(
    key, ensamble_size
)

params, train_loss, test_loss = vmap(train_network)(ensemble_key)

perdict_fn = nt.predict.gradient_descent_mse_ensemble(
    kernel_fn, train_xs, train_ys, diag_reg=1e-4
)
ntk_mean, ntk_var = perdict_fn(
    x_test=test_xs, get="ntk", compute_cov=True
)

ntk_mean = np.reshape(ntk_mean, (-1,))
ntk_std = np.sqrt(np.diag(
    ntk_var
))

ntk_train_loss_mean = loss_fn(perdict_fn, train_ys, ts)
ntk_test_loss_mean = loss_fn(perdict_fn, test_ys, ts, test_xs)

plt.subplot(1, 2, 1)

mean_train_loss = np.mean(train_loss, axis=0)
mean_test_loss = np.mean(test_loss, axis=0)

plt.loglog(ts, ntk_train_loss_mean, linewidth=3)
plt.loglog(ts, ntk_test_loss_mean, linewidth=3)

plt.loglog(ts, mean_train_loss, 'k-', linewidth=2)
plt.loglog(ts, mean_test_loss, 'k-', linewidth=2)

plt.xlim([10 ** 0, 10 ** 3])

plt.xscale('log')
plt.yscale('log')

plt.legend(['Infinite Train', 'Infinite Test', 'Finite Ensemble'])

plt.subplot(1, 2, 2)

plot_fn(train, None)

plt.plot(test_xs, ntk_mean, 'b-', linewidth=3)
plt.fill_between(
    np.reshape(test_xs, (-1)),
    ntk_mean - 2 * ntk_std,
    ntk_mean + 2 * ntk_std,
    color='blue', alpha=0.2)

ensemble_fx = vmap(apply_fn, (0, None))(params, test_xs)

mean_fx = np.reshape(np.mean(ensemble_fx, axis=0), (-1,))
std_fx = np.reshape(np.std(ensemble_fx, axis=0), (-1,))

plt.plot(test_xs, mean_fx - 2 * std_fx, 'k--', label='_nolegend_')
plt.plot(test_xs, mean_fx + 2 * std_fx, 'k--', label='_nolegend_')
plt.plot(test_xs, mean_fx, linewidth=2, color='black')

plt.legend(['Train', 'Infinite Network', 'Finite Ensemble'], loc='upper left')

plt.xlim([-np.pi, np.pi])
plt.ylim([-1.5, 1.5])
plt.savefig("./plot/resnet")


import pandas as pd
import matplotlib.pyplot as plt

# Read the spreadsheet
df = pd.read_csv("downloads.csv", header=None, names=["Date", "Downloads"])
df["Date"] = pd.to_datetime(df["Date"])

# Plot the download data
plt.figure(figsize=(10, 6))
plt.plot(df["Date"], df["Downloads"])
plt.title("Downloads Over Time")
plt.xlabel("Date")
plt.ylabel("Downloads")
plt.savefig("downloads.png")


"""
Module containing the SinusoidalPositionalEmbedding option, which creates a sinusoidal relationship
between a token and its position
"""

import logging
from rich.logging import RichHandler

LOG_FORMAT = "%(message)s"
logging.basicConfig(level="INFO", format=LOG_FORMAT, datefmt="[%X] ", handlers=[RichHandler()])
LOGGER = logging.getLogger(__name__)

import os
import sys

import math

import torch
from torch import nn
import torch.nn.functional as F
import torch.onnx.operators

from annotated_mpnet.utils import utils

class SinusoidalPositionalEmbedding(nn.Module):
    """
    A module for creating positional embeddings that follow a sinusoidal relationship
    """

    def __init__(self, embedding_dim, padding_idx, init_size=1024) -> None:
        super().__init__()

        # Store args
        self.embedding_dim = embedding_dim
        self.padding_idx = padding_idx

        # Get the weights from a helper function that processes the sinusoidal math
        self.weights = SinusoidalPositionalEmbedding.get_embedding(
            init_size,
            embedding_dim,
            padding_idx,
        )

        # Set the ONNX trace variable again, but I don't think we'll be using it
        self.onnx_trace = False

        # This is a builtin nn.Module function for registering buffer values within a module
        self.register_buffer("_float_tensor", torch.FloatTensor(1))

    def prepare_for_onnx_export(self):
        """
        This function is probably useless for us, but we keep it in
        """
        self.onnx_trace = True

    @staticmethod
    def get_embedding(num_embeddings, embedding_dim, padding_idx=None):
        """
        This static method instantiates the sinusoidal pattern for the positional embeddings
        """

        # First get half the embedding dimension size
        half_dim = embedding_dim // 2

        # Not quite sure of the math happening below, but generally, we are constructing initial
        # weights using an algorithm that heavily involves trigonometric relationships (as you can
        # see in the last step with sin() and cos() making an appearance)
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb)
        emb = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(1) * emb.unsqueeze(0)
        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(num_embeddings, -1)

        # Next calculate padding. If embedding size is not divisible by 2, we need to pad out
        if embedding_dim % 2 == 1:
            # zero pad
            emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)

        # If there IS a padding index, reset the weights to 0
        if padding_idx is not None:
            emb[padding_idx, :] = 0
        return emb

    def forward(self, input, incremental_state=None, timestep=None, **kwargs):
        """
        The forward function for processing these positional embeddings. Input should be of size
        (bsz x seq_len) as usual
        """

        # Break out dimensions
        bsz, seq_len = torch.onnx.operators.shape_as_tensor(input)

        # Get the max position of the given sequence
        max_pos = (self.padding_idx + 1) + seq_len

        # Now we add the option to recompute embeddings if the initial embeddings weren't large
        # enough to cover all positons
        if self.weights is None or max_pos > self.weights.size(0):
            # recompute/expand embeddings if needed
            self.weights = SinusoidalPositionalEmbedding.get_embedding(
                max_pos,
                self.embedding_dim,
                self.padding_idx,
            )

        self.weights = self.weights.to(self._float_tensor)

        # Process incremental state below
        # Again, not really sure what this does
        if incremental_state is not None:
            # positions is the same for every token when decoding a single step
            pos = timestep.view(-1)[0] + 1 if timestep is not None else seq_len
            if self.onnx_trace:
                return (
                    self.weights.index_select(index=self.padding_idx + pos, dim=0)
                    .unsqueeze(1)
                    .repeat(bsz, 1, 1)
                )
            return self.weights[self.padding_idx + pos, :].expand(bsz, 1, -1)

        # Use the typical `make_positions` util to get incremental positions. This will eventually
        # feed directly into the sinusoidal weights we defined before
        positions = utils.make_positions(input, self.padding_idx, onnx_trace=self.onnx_trace)

        # If onnx_trace is set (which it shouldn't be), process additional below
        if self.onnx_trace:
            flat_embeddings = self.weights.detach().index_select(0, positions.view(-1))
            embedding_shape = torch.cat((bsz.view(1), seq_len.view(1), torch.LongTensor([-1])))
            embeddings = torch.onnx.operators.reshape_from_tensor_shape(
                flat_embeddings, embedding_shape
            )
            return embeddings

        # Return the weights selected by the positions generated above
        return self.weights.index_select(0, positions.view(-1)).view(bsz, seq_len, -1).detach()

    # Helper function below
    def max_positions(self):
        """
        Maximum number of supported positions.
        """
        return int(1e5)  # an arbitrary large number


from app.constants import SEASON
from app.enums import CharacterClass, Server
from app.extensions import db
from app.models import (FinishedMatch, FinishedMatchParticipant, Player,
                        PlayerStatistics, RunningMatch)
from flask import Blueprint, abort, request, jsonify
from flask.globals import current_app
from flask.helpers import send_from_directory
from flask.templating import render_template
from flask_babel import gettext
from sqlalchemy import or_, text

bp = Blueprint("main", __name__)

@bp.route("/", methods=["GET"])
def index():
    return render_template("main/index.html",
                           title=gettext("Home"))

@bp.route("/robots.txt", methods=["GET"])
def robots():
    return send_from_directory(current_app.static_folder, "robots.txt")

@bp.route("/sitemap.xml", methods=["GET"])
def sitemap():
    return send_from_directory(current_app.static_folder, "sitemap.xml")

@bp.route("/privacy-policy", methods=["GET"])
def privacy_policy():
    return render_template("main/privacy_policy.html",
                           title=gettext("Privacy Policy"))

@bp.route("/legal-notice", methods=["GET"])
def legal_notice():
    return render_template("main/legal_notice.html",
                           title=gettext("Legal Notice"))

@bp.route("/ranking/<server>", methods=["GET"])
def ranking(server: str):
    if server not in ["bergruen", "luxplena"]:
        abort(404)

    season = request.args.get("season", SEASON)
    server: Server = Server(server)
    if server == server.bergruen:
        title = gettext("Ranking Bergruen")
    else:
        title = gettext("Ranking LuxPlena")

    # Query players
    query = (
        db.session.query(
            Player, PlayerStatistics
        ).join(
            PlayerStatistics
        ).filter(
            Player.exists,
            Player.server == server,
            PlayerStatistics.season == season,
            or_(
                PlayerStatistics.wins > 0,
                PlayerStatistics.losses > 0,
                PlayerStatistics.draws > 0,
            ),
        ).order_by(
            PlayerStatistics.points.desc(),
            Player.username.desc(),
        )
    )
    players = query.all()

    return render_template(
        "main/ranking.html",
        title=title,
        server=server.value,
        players=players,
    )

@bp.route("/matches/<server>/", methods=["GET"])
def matches(server: str):
    if server not in ["bergruen", "luxplena"]:
        abort(404)

    server: Server = Server(server)
    if server == server.bergruen:
        title = gettext("Matches Bergruen")
    else:
        title = gettext("Matches LuxPlena")

    matches = (
        db.session.query(
            RunningMatch
        ).filter(
            RunningMatch.server == server,
        ).all()
    )

    # Check if a json flag is set
    if int(request.args.get("json", 0)) == 1:
        json_matches = []
        for match in matches:
            match: RunningMatch

            json_matches.append({
                "id": match.id,
                "server": server.value,
                "team_1": [p.player.to_dict() for p in match.team_1],
                "team_2": [p.player.to_dict() for p in match.team_2],
            })

        return jsonify(json_matches), 200

    return render_template(
        "main/matches.html",
        title=title,
        server=server.value,
        matches=matches,
    )

@bp.route("/players/<server>/<int:id>")
def player_profile(server: str, id: int):
    server: Server = Server(server)
    season = request.args.get("season", SEASON)

    query_res = (
        db.session.query(
            Player, PlayerStatistics
        ).join(
            PlayerStatistics
        ).filter(
            Player.exists,
            Player.id == id,
            PlayerStatistics.season == season,
        ).first()
    )

    if not query_res:
        abort(404, "Player not found")

    player, statistics = query_res

    matches = (
        db.session.query(
            FinishedMatch,
        ).join(
            FinishedMatchParticipant
        ).filter(
            FinishedMatch.season == season,
            FinishedMatchParticipant.player_id == player.id,
        ).order_by(
            FinishedMatch.date.desc()
        ).all()
    )

    return render_template(
        "main/player_profile.html",
        player=player,
        statistics=statistics,
        matches=matches,
        title=gettext("Player %(name)s", name=player.username),
    )

@bp.route("/<server>/winners")
def server_winners(server: str):
    server: Server = Server(server)
    season = request.args.get("season", SEASON)
    matches_required = 16

    classes = {
        "noble": [CharacterClass.noble, CharacterClass.court_magician, CharacterClass.magic_knight],
        "saint": [CharacterClass.saint, CharacterClass.shaman, CharacterClass.priest],
        "explorer": [CharacterClass.explorer, CharacterClass.sniper, CharacterClass.excavator],
        "mercenary": [CharacterClass.mercenary, CharacterClass.gladiator, CharacterClass.guardian_swordsman],
    }

    class_based_data = {
        "noble": [],
        "saint": [],
        "explorer": [],
        "mercenary": [],
    }

    for key, key_classes in classes.items():
        query = (
            db.session.query(
                Player, PlayerStatistics,
                (PlayerStatistics.wins + PlayerStatistics.losses + PlayerStatistics.draws).label("matches_played"),
            ).join(
                PlayerStatistics
            ).filter(
                Player.exists,
                Player.server == server,
                PlayerStatistics.season == season,
                text(f"matches_played >= {matches_required}"),
                Player.character_class.in_(key_classes)
            ).order_by(
                PlayerStatistics.points.desc(),
                Player.username.desc(),
            ).limit(10)
        )

        class_based_data[key].extend(query.all())

    overall_data = (
        db.session.query(
            Player, PlayerStatistics,
            (PlayerStatistics.wins + PlayerStatistics.losses + PlayerStatistics.draws).label("matches_played"),
        ).join(
            PlayerStatistics
        ).filter(
            Player.exists,
            Player.server == server,
            PlayerStatistics.season == season,
            text(f"matches_played >= {matches_required}"),
        ).order_by(
            PlayerStatistics.points.desc(),
            Player.username.desc(),
        ).limit(10)
    ).all()

    prizes = {
        "overall": {
            1: "25.000 AP",
            2: "20.000 AP",
            3: "15.000 AP",
            4: "10.000 AP",
            5: " 5.000 AP",
        },
        "class_based": {
            1: "15x Slate Piece of Ancient (uh0000008)",
            2: "10x Slate Piece of Ancient (uh0000008)",
            3: "5x Slate Piece of Ancient (uh0000008)",
        }
    }

    return render_template("main/winners.html",
                           class_based_data=class_based_data,
                           overall_data=overall_data,
                           prizes=prizes,
                           matches_required=matches_required)


def pairwise_swap1(n):
    mask = 0b11
    count = 0
    c = n
    while c != 0:
        if c & mask == 1 or c & mask == 2:
            n ^= mask << count
        c = rshift(c, 2)
        count += 2
    return n

def rshift(a, b):
    return (a % 0x100000000) >> b

# 5 instructions only
# assume 32-bit integer
def pairwise_swap2(n):
    odd_mask = 0xaaaaaaaa
    even_mask = 0x55555555
    return rshift(n & odd_mask, 1) | (n & even_mask) << 1

cases = [0b1001101, 0, 1, 0b111, 0b1111, 0b110]
for i, case in enumerate(cases):
    print(f"case {i+1}: {case:b}")
    print(f"{pairwise_swap1(case):b}")
    print(f"{pairwise_swap2(case):b}")


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Oct 14 11:51:09 2020

@author: apurvabadithela
"""
import networkx as nx
from networkx.algorithms.flow import shortest_augmenting_path, edmonds_karp
import sys
sys.path.append('../')

from src.milp_functions import augment_paths, min_cut_edges
from src.restrict_transitions_cycles import remove_edges_corrected_cycle, remove_edges
from src.simulation_helpers import post_process_cuts, process_history

# Constructing graph induced by automaton:
G = nx.DiGraph()
G.add_edge("n1", "n2", capacity=1.0)
G.add_edge("n2", "n3", capacity=1.0)
G.add_edge("n2", "n4", capacity=1.0)
G.add_edge("n2", "n5", capacity=1.0)
G.add_edge("n3", "n6", capacity=1.0)
G.add_edge("n4", "n6", capacity=1.0)
G.add_edge("n5", "n6", capacity=1.0)
G.add_edge("n6", "n7", capacity=1.0)
G.add_edge("n2", "n8", capacity=1.0)
G.add_edge("n8", "n5", capacity=1.0)

# Specifying the goal and propositions:
g ="n7"
w = "n3"
q0 = "n1"

P = augment_paths(G, [q0], [g])
MC = min_cut_edges(G, [q0], [g])

props = [[q0], [w], [g]]
C, Q0, Chist, discard, nkeep, niterations, nf_iterations, alg_fail_main, skip_iter = remove_edges(G, props, props[0], "ALL")

C = post_process_cuts(C, G)
Chist = process_history(Chist,G)
print("List of constrained edges: ")
print(C)
print("List of constrained edges in each iteration: ")
print(Chist)


import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torchvision import datasets, transforms
from utils import weights_init_xavierUniform as weights_init
import math

class Flatten(nn.Module):
    def forward(self, input):
        return input.view(input.size(0), -1)

class Recurrent_block_old(nn.Module):
    def __init__(self,ch_out,t=2):
        super(Recurrent_block,self).__init__()
        self.t = t
        self.ch_out = ch_out
        self.conv = nn.Sequential(
            nn.Conv2d(ch_out,ch_out,kernel_size=3,stride=1,padding=1,bias=True),
		    nn.BatchNorm2d(ch_out),
			nn.ReLU(inplace=True)
        )

    def forward(self,x):
        for i in range(self.t):

            if i==0:
                x1 = self.conv(x)

            x1 = self.conv(x+x1)
        return x1

class Recurrent_block(nn.Module):
    def __init__(self,ch_out,t=2):
        super(Recurrent_block,self).__init__()
        self.t = t
        self.conv = nn.Conv2d(ch_out,ch_out,kernel_size=3,stride=1,padding=1,bias=True)
        self.bn = nn.ModuleList([nn.BatchNorm2d(ch_out) for i in range(t)])
        self.relu = nn.ReLU(inplace=True)

        self.shortcut = nn.Conv2d(ch_out,ch_out,kernel_size=3,stride=1,padding=1,bias=True)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.xavier_uniform_(m.weight.data)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)

    def forward(self,x):
        rx = x
        for i in range(self.t):

            if i==0:
                x1 = self.conv(x)

            x1 = self.conv(x) + self.shortcut(rx)
            x = self.relu(x1)
            x = self.bn[i](x)
        return x

class Attention_block(nn.Module):
    def __init__(self,F_g,F_l,F_int):
        super(Attention_block,self).__init__()
        self.W_g = nn.Sequential(
            nn.Conv2d(F_g, F_int, kernel_size=1,stride=1,padding=0,bias=True),
            nn.BatchNorm2d(F_int)
            )

        self.W_x = nn.Sequential(
            nn.Conv2d(F_l, F_int, kernel_size=1,stride=1,padding=0,bias=True),
            nn.BatchNorm2d(F_int)
        )

        self.psi = nn.Sequential(
            nn.Conv2d(F_int, 1, kernel_size=1,stride=1,padding=0,bias=True),
            nn.BatchNorm2d(1),
            nn.Sigmoid()
        )

        self.relu = nn.ReLU(inplace=True)

    def forward(self,g,x):
        g1 = self.W_g(g)
        x1 = self.W_x(x)
        psi = self.relu(g1+x1)
        psi = self.psi(psi)

        return x*psi

class Conv_block_2(nn.Module):
    def __init__(self, ch_in, ch_mid, ch_out, k1, k2, s1, s2, p1, p2):
        super(Conv_block_2, self).__init__()
        self.conv_1 = nn.Sequential(
            nn.Conv2d(ch_in, ch_mid, kernel_size=k1, stride=s1, padding=p1, bias=True),
            nn.ReLU(inplace=True)
        )
        self.conv_2 = nn.Sequential(
            nn.Conv2d(ch_mid, ch_out, kernel_size=k2, stride=s2, padding=p2, bias=True),
            nn.ReLU(inplace=True)
        )
        self.bn = nn.BatchNorm2d(ch_mid, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        self.apply(weights_init)

    def forward(self, x):
        x1 = self.conv_1(x)
        x1 = self.bn(x1)
        x2 = self.conv_2(x1)
        # print('Working now')
        # print('x2',x2)
        return x2

class Conv_block_3(nn.Module):
    def __init__(self, ch_in, ch_mid, ch_mid_2, ch_out, k1, k2, k3, s1, s2, s3, p1, p2, p3):
        super(Conv_block_3, self).__init__()
        self.conv_1 = nn.Sequential(
            nn.Conv2d(ch_in, ch_mid, kernel_size=k1, stride=s1, padding=p1, bias=True),
            nn.ReLU(inplace=True),
            nn.Conv2d(ch_mid, ch_mid_2, kernel_size=k2, stride=s2, padding=p2, bias=True),
            nn.ReLU(inplace=True)
        )
        self.conv_2 = nn.Sequential(
            nn.Conv2d(ch_mid_2, ch_out, kernel_size=k3, stride=s3, padding=p3, bias=True),
            nn.ReLU(inplace=True)
        )
        self.bn = nn.BatchNorm2d(ch_mid_2, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        self.apply(weights_init)

    def forward(self, x):
        x1 = self.conv_1(x)
        x2 = self.bn(x1)
        x3 = self.conv_2(x2)
        # print('Working now')
        # print('x3',x3)
        return x3


T=input()
for t in range(int(T)):
	x,y,TT=map(int, input().split())
	print(x,y,TT)
	arr=[]
	arr.append(['0' for i in range(y+2)])
	for i in range(x):
		tmp=[]
		tmp+=list(input())
		tmp.append('0')
		tmp.insert(0,'0')
		arr.append(tmp)
	arr.append(['0' for i in range(y+2)])
	for tt in range(TT):
		a,b=map(int, input().split())
		a=a+1
		b=b+1
		count=1
		p=0
		for i in range(1,100):
			for j in range(a-i,a+i+1):
				for k in range(b-i,b+i+1):
				#	print(j,k)
					if arr[a][b]!=arr[j][k]:
						p=1
			if p==1:
				break
			count+=2
		print(count)
# T=input()
# for t in range(int(T)):
# 	arr=list(map(int, input().split()))
# 	print(arr[0],arr[1],arr[2])
# 	squ=[]
# 	for i in range(arr[0]):
# 		squ.append(list(input()))
# 	for i in range(int(arr[2])):
# 		a,b=map(int, input().split())
# 		count=1
# 		p=0
# 		while p==0:
# 			for n in range(1,100):
# 				for j in range(a-n,a+n+1):
# 					for k in range(b-n,b+n+1):
# 						if 0<=j<arr[0] and 0<=k<arr[1]:
# 							if squ[j][k]!=squ[a][b]:
# 								p=1
# 						else:
# 							p=1
# 				if p==0:
# 					count+=2
# 				else:
# 					break
# 		print(count)


#!/usr/bin/env python3

from argparse import ArgumentParser
from functools import reduce
import json
import os

def add(a, b):
    return a + b

def get_file_statistics(file):
    with open(file, 'r') as f:
        content = json.loads(f.read())

        return {
            'filename': os.path.basename(file),
            'repo': content['repo'],
            'streaks': len(content['streaks']),
            'commits': reduce(add, [len(streak['commits']) for streak in content['streaks']]),
        }

def main():
    parser = ArgumentParser()
    parser.add_argument('dir', help='Directory to generate index for')
    args = parser.parse_args()

    json_dir = os.path.realpath(args.dir)
    indexed_files = []

    for file in os.listdir(json_dir):
        if file == '_index.json':
            continue

        print('Parsing {}'.format(file))
        file_path = os.path.join(json_dir, file)
        indexed_files.append(get_file_statistics(file_path))

    with open(os.path.join(json_dir, '_index.json'), 'w') as f:
        f.write(json.dumps(indexed_files, indent=2))

if __name__ == '__main__':
    main()


class Solution:
    def getSeq(self, nums, i, n):
        if i >= n:
            return [[]]
        op = []
        used = set()
        for j in range(i + 1, n):
            if nums[j] >= nums[i] and nums[j] not in used:
                val = self.getSeq(nums, j, n)
                op += [[nums[i], nums[j]]]
                op += [[nums[i]] + v for v in val]
                used.add(nums[j])
        return op

    def findSubsequences(self, nums: List[int]) -> List[List[int]]:
        n = len(nums)
        op = []
        used = set()
        for i in range(n):
            if nums[i] not in used:
                op += self.getSeq(nums, i, n)
                used.add(nums[i])
        return op


import socket
import threading
import argparse

def receive_messages(client_socket):
    while True:
        try:
            message = client_socket.recv(1024).decode('utf-8')
            print(message)
        except ConnectionResetError:
            print("Disconnected from the server.")
            break

def main():
    parser = argparse.ArgumentParser(description="GoChat Client")
    parser.add_argument(
        "-s", "--server", default="localhost", help="Server address")
    parser.add_argument("-p", "--port", type=int,
                        default=6969, help="Port number")
    args = parser.parse_args()

    host = args.server
    port = args.port

    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    client_socket.connect((host, port))

    receive_thread = threading.Thread(
        target=receive_messages, args=(client_socket,))
    receive_thread.daemon = True
    receive_thread.start()

    while True:
        message = input()
        if message == "exit":
            break
        client_socket.send(message.encode('utf-8'))

if __name__ == "__main__":
    main()


import numpy as np
from numpy.random.mtrand import f
from scipy.spatial.transform import Rotation as R

# function for the identification of closest center of atoms within a 7Å radius
def find_closest_centroid(mol, i_point, rad=7):
    neigh = []
    for e in mol:
        rf = np.linalg.norm(e-i_point)
        if rf < rad:
            neigh.append([e[0],e[1],e[2]])

    neighbours = np.array([np.array(x) for x in neigh])
    close_centroid = np.array((sum(neighbours[:,0])/len(neighbours[:,0]),sum(neighbours[:,1])/len(neighbours[:,1]),sum(neighbours[:,2])/len(neighbours[:,2])))

    return close_centroid

def solvent_placer(solv_repo, ssip, solute_cent):
    keys = [l[0] for l in solv_repo]
    solvent_array = np.array([l[1] for l in solv_repo])

    vo = solvent_array[-2]
    xo = ssip
    solvent_o = np.array([at-vo for at in solvent_array]) # translate standard solvent coordinates at the origin

    ox, oy, oz = solute_cent
    xx, xy, xz = xo - solute_cent
    Vx, Vy, Vz = solvent_o[-2]
    bx, by, bz = solvent_o[-1]
    r = np.array([xx, xy, xz])
    s = np.array([bx-Vx, by-Vy, bz-Vz])
    rs = np.cross(r,s)

    rs2 = rs / np.sqrt(np.dot(rs, rs))

    alfa = np.arccos((np.dot(r,s)/(np.linalg.norm(r) * np.linalg.norm(s))))

    # translation + first rotation

    r1 = R.from_rotvec(alfa * rs2)
    rot1 = r1.apply(solvent_o)

    if np.allclose(r, rot1[-1], atol=1e-01) == False:
        r_neg = R.from_rotvec(-alfa * rs2)
        rot1 = r_neg.apply(solvent_o)
    else:
        pass

    betas = np.arange(0,3.14, 0.5)

    rotamer = []
    distance = 0
    for b in betas:
        ac = rot1[-1] - rot1[-2]
        l = ssip - rot1[-2]
        r2 = R.from_rotvec(b * ac)
        rot2 = r2.apply(rot1)

        finalfinal = np.array([at + l for at in rot2])

        fdist = [np.linalg.norm(k - solute_cent) for k in finalfinal[0:-2]]
        avg_dist = sum(fdist)/len(fdist)

        if avg_dist > distance:
            distance = avg_dist
            rotamer = finalfinal
        else:
            pass

    ex1 = [[str(t) for t in l] for l in rotamer]
    for i in range(len(ex1)):
        rot_atom = ex1[i]
        rot_atom.insert(0, keys[i])

    solv = [e for e in ex1 if e[0] not in ('Cnt', 'Vx')]
    return solv

def cluster_maker(solute, int_points, sv_dict, sv_name):
    hybrid_solute = [[i for i in d] for d in solute]
    for key in int_points:
        esp = key[3] # esp value of interaction points
        ssip = np.array(key[0:3]) # coordinates of interaction points

        #create numpy array of solute coordinates and localize the centeroid for optimal solvent orientation
        solute_coord = [[float(c) for c in i[1:]] for i in hybrid_solute]
        at_array = np.array([np.array(x) for x in solute_coord])
        closest = find_closest_centroid(at_array, ssip) # finds centroid in a 5A radius from the inter. point

        # place solvent molecules according to ESP value of the interaction point and centroid position.
        # Newly placed solvent molecule is added to solute atoms array and its atoms are considered in
        # the identification of the new centroid for the successive solvent placement iteration
        if esp > 0.0001:
            solvent_mode = sv_dict[sv_name + '_min']
            exp_solv = solvent_placer(solvent_mode, ssip, closest)
            for q in exp_solv:
                hybrid_solute.append(q)

        elif esp < -0.0001:
            solvent_mode = sv_dict[sv_name + '_max']
            exp_solv = solvent_placer(solvent_mode, ssip, closest)
            for q in exp_solv:
                hybrid_solute.append(q)

        else:
            pass

    hybrid_solute

    return hybrid_solute


from django.db import models
from tumblr.base_model import BaseManager, BaseModel
from apps.blog.models import Post

class TagsManager(BaseManager):

  def posts_with_tags(self, tags, post_count, user=False):
    end_count = post_count + 10

    tag_list = []
    split_tags = tags.replace('+', ' ').split(' ')

    for tag in split_tags:

      if tag == '':
        pass
      else:
        tagss = self.get(tags=tag)
        tag_list.append(tagss.pk)

    if user:
      post_with_tags = Post.objects.exclude(user=user)
    else:
      post_with_tags = Post.objects.all()

    for pk in tag_list:
      loop_posts = post_with_tags.filter(tags__pk=pk)
      post_with_tags = loop_posts

    ordered_posts = post_with_tags.order_by('-pub_date')[post_count:end_count]

    return ordered_posts

  def filter_by_post(self, post):
    return self.filter(post=post)

  def create_tags(self, new_tags, post):
    stripped_tags = new_tags.replace('#', '').strip()
    split_tags = stripped_tags.split(' ')

    for tag in split_tags:

      if tag == '':
        pass
      else:

        try:
          tags = Tags.objects.get(tags=tag)
          tags.post.add(post)
        except:
          tags = Tags(tags=tag)
          tags.save()
          tags.post.add(post)

class Tags(BaseModel):

  tags = models.CharField(max_length=20, unique=True, blank=True)
  post = models.ManyToManyField(Post)
  objects = TagsManager()

  class Meta:
    app_label = 'search'

  def __str__(self):
    return self.tags