#!/usr/bin/env python3

import asyncio
import random
import sys

resource = ['*', '$', '#']
resource_lock = asyncio.Lock()
#排队锁，不论执行读操作还是写操作都要先获得该锁，在获得资源锁后释放。
queue_lock = asyncio.Lock()
read_count = 0

async def writer(data, name):
    global resource, resource_lock, queue_lock
    #申请获得排队锁。
    await queue_lock.acquire()
    await resource_lock.acquire()
    #释放排队锁。
    queue_lock.release()
    for i in range(3):
        t = random.random()
        await asyncio.sleep(t)
        resource[i] = data[i]
        print(f"{name}: write {data[i]}, now {resource[0] + resource[1] + resource[2]}")
    resource_lock.release()

async def reader(name):
    global resource, resource_lock, read_count, queue_lock
    #申请获得排队锁。
    await queue_lock.acquire()
    read_count = read_count + 1
    if read_count == 1:
        await resource_lock.acquire()
    #释放排队锁。
    queue_lock.release()
    for i in range(3):
        t = random.random()
        await asyncio.sleep(t)
        print(f"{name}: read {resource[i]}")
    read_count = read_count - 1
    if read_count == 0:
        resource_lock.release()

async def main():
    r1 = asyncio.create_task(reader("r1"))
    w1 = asyncio.create_task(writer("abc", "w1"))
    r2 = asyncio.create_task(reader("r2"))
    r3 = asyncio.create_task(reader("r3"))
    w2 = asyncio.create_task(writer("def", "w2"))
    r4 = asyncio.create_task(reader("r4"))
    await asyncio.gather(r1, w1, r2, r3, w2, r4)
    sys.exit(0)

if __name__ == '__main__':
    asyncio.run(main())


# day 3 was done on day 2

rock = '''
    _______
---'   ____)
      (_____)
      (_____)
      (____)
---.__(___)
'''

paper = '''
    _______
---'   ____)____
          ______)
          _______)
         _______)
---.__________)
'''

scissors = '''
    _______
---'   ____)____
          ______)
       __________)
      (____)
---.__(___)
'''
import random
def rps():
    print("ROCK PAPER SCISSORS\nYou vs Computer\nEnter 'q' to quit\n\n")
    choices =[rock, paper, scissors]
    choice_player= int(input("What do you choose? Type 0 for Rock, 1 for Paper or 2 for Scissors."))
    computer_choice = random.choice(choices)

    winner = {rock:scissors, paper:rock, scissors:paper}
    print(f'You choose:\n{choices[choice_player]}\nComputer choose:\n{computer_choice}')
    choice_player = choices[choice_player]

    if winner[choice_player] == computer_choice:
        return print('You Won')
    elif winner[computer_choice] == choice_player:
        return print('You lose')
    else:
        return print('Draw')

rps()


"""Module for interacting with MATL, and it's source code."""
import pathlib
from typing import Optional

from matl_online.octave import OctaveSession, OutputCallback
from matl_online.octave import string as octave_string
from matl_online.types import MATLTaskParameters

from .source import get_matl_folder

def matl(
    octave: OctaveSession,
    matl_params: MATLTaskParameters,
    directory: pathlib.Path,
    line_handler: Optional[OutputCallback] = None,
) -> None:
    """Open a session with Octave and manages input/output as well as errors."""

    # Add the folder for the appropriate MATL version
    matl_folder = get_matl_folder(matl_params.version)

    # Ensure the matl folder exists
    assert matl_folder, "MATL folder does not exist"

    # Change directories to the temporary folder so that all temporary
    # files are placed in here and won't interfere with other requests
    with octave.current_directory(directory):
        with octave.paths(matl_folder):
            # Convert the code to a cell array element-per-line
            code = f"{{{','.join([octave_string(x) for x in matl_params.code_lines])}}}"

            octave.run(
                "matl_runner",
                octave_string(matl_params.flags),
                code,
                *[octave_string(x) for x in matl_params.input_lines],
                line_handler=line_handler,
            )


import pygame
import sys
import random
import os
os.chdir("C:\\Users\\박예은\\AppData\\Local\\Programs\\Python\\Python37\\One_Big_Archive\\images")

class MyBallClass(pygame.sprite.Sprite):
    def __init__(self, image_file, location, speed):
        pygame.sprite.Sprite.__init__(self)     # call Sprite initializer ; 스프라이트 초기화
        self.image = pygame.image.load(image_file)
        self.rect = self.image.get_rect()
        self.rect.left, self.rect.top = location
        self.speed = speed

    def move(self):
        self.rect = self.rect.move(self.speed)
        if self.rect.left < 0 or self.rect.left > width:
            self.speed[0] = -self.speed[0]

        if self.rect.top < 0 or self.rect.bottom > height:
            self.speed[1] = -self.speed[1]

size = width, height = 640, 480
screen = pygame.display.set_mode(size)
# screen = pygame.display.set_mode([640, 480])
# print(width, type(width))
# print(size, type(size))
screen.fill([255, 255, 255])
img_file = "beach_ball.png"
balls = []
for row in range(0, 3):
    for col in range(0, 3):
        location = [col*180+10, row*180+10]
        speed = [random.choice([-2, 2]), random.choice([-2, 2])]
        ball = MyBallClass(img_file, location, speed)
        balls.append(ball)

while True:
    for i in pygame.event.get():
        if i.type == pygame.QUIT:
            sys.exit()
    pygame.time.delay(20)
    screen.fill([255, 255, 255])
    for ball in balls:
        ball.move()
        screen.blit(ball.image, ball.rect)
    pygame.display.flip()


import pandas as pd
import numpy as np

df = pd.read_csv(r'Y:\project\data\data_tmc2.csv')
#df = pd.read_csv('/Users/yuanjielu/Desktop/python/project/data/data_tmc.csv')
df_tmc = pd.read_csv(r'Y:\project\data\tmc.csv')
#df_tmc = pd.read_csv('/Users/yuanjielu/Desktop/python/project/data/tmc.csv')
time = pd.read_csv(r'Y:\project\data\time_random.csv')
#time = pd.read_csv('/Users/yuanjielu/Desktop/python/project/data/time_random.csv')
tmc_list = list(df_tmc.tmc)
time_list = list(time.time)
#time = df['measurement_tstamp'].drop_duplicates()
#time = pd.DataFrame(time)
#time.to_csv("X:/project/data/time_random.csv")
print(time)
df.measurement_tstamp = pd.to_datetime(df.measurement_tstamp)
time.time = pd.to_datetime(time.time)
speed_adj = pd.DataFrame(np.zeros(shape = (len(time_list),len(tmc_list))), index = time_list, columns = tmc_list)
construction_adj = pd.DataFrame(np.zeros(shape = (len(time_list),len(tmc_list))), index = time_list, columns = tmc_list)
for i in range(len(time_list)):
    tmc = df[df['measurement_tstamp'] == time['time'].iloc[i]]
    for j in range(len(tmc)):
        b_index = tmc_list.index(tmc['tmc_code'].iloc[j])
        print(tmc['speed'].iloc[j])
        speed_adj.iat[i,b_index ] = tmc['speed'].iloc[j]
        construction_adj.iat[i, b_index] = tmc['Work_zone'].iloc[j]
speed_adj.to_csv(r'Z:\project\data\adj\new_speed_adj_10.csv', index = True, header=True)
construction_adj.to_csv(r'Z:\project\data\adj\new_construction_adj_10.csv', index = True, header=True)
#speed_adj.to_csv('/Users/yuanjielu/Desktop/python/project/data/adj/speed_adj.csv', index = True, header = True)
#construction_adj.to_csv('/Users/yuanjielu/Desktop/python/project/data/adj/construction_adj_10.csv', index = True, header = True)


from keras.preprocessing.image import ImageDataGenerator, array_to_img, img_to_array, load_img
from matplotlib.pyplot import imread, imshow, show, savefig
import numpy as np
import sys
import os
import cv2

WALDO_DIR = os.path.dirname(os.path.realpath(__file__)) + '/../Hey-Waldo/'
WALDO_DUPLICATION_NB = 100
NOTWALDO_DUPLICATION_NB = 2
NOISE_DUPLICATION = 30

def get_images_data(dir):
    return np.array([np.array(imread(WALDO_DIR + dir + '/waldo/' + fname))
                     for fname in os.listdir(WALDO_DIR + dir + '/waldo')]),\
            np.array([np.array(imread(WALDO_DIR + dir + '/notwaldo/' + fname))
                      for fname in os.listdir(WALDO_DIR + dir + '/notwaldo')])

def data_augmentation(waldos, notWaldos, dir):
    datagen = ImageDataGenerator(
        rotation_range=40,
        width_shift_range=0.2,
        height_shift_range=0.2,
        rescale=1. / 255,
        horizontal_flip=True,
        fill_mode='nearest')

    print('Waldo data augmentation...')
    for waldo in waldos:
        waldo = waldo.reshape((1,) + waldo.shape)

        i = 0
        for batch in datagen.flow(waldo, batch_size=1,
                                  save_to_dir=WALDO_DIR + dir + '/waldoOther', save_prefix='waldo', save_format='jpeg'):
            i += 1
            if i >= WALDO_DUPLICATION_NB - 1:
                break

    print('Not waldo data augmentation...')
    for not_waldo in notWaldos:
        not_waldo = not_waldo.reshape((1,) + not_waldo.shape)

        i = 0
        for batch in datagen.flow(not_waldo, batch_size=1,
                                  save_to_dir=WALDO_DIR + dir + '/notWaldoOther', save_prefix='waldo', save_format='jpeg'):
            i += 1
            if i >= NOTWALDO_DUPLICATION_NB - 1:
                break

def get_waldos_data(path):
    waldos = []
    for fname in os.listdir(path + '/'):
        waldos.append(cv2.imread(path + '/' + fname))
    return np.array(waldos)

def data_add_noise(waldos, path):
    for idx, waldo in enumerate(waldos):
        for i in range(40):
            noise = (0.4 * np.random.randn(*waldo.shape)) * 10
            img = waldo + noise
            cv2.imwrite(path + '/noisy_' + str(idx) + '_' + str(i) + '.jpg', img)

if __name__ == '__main__':
    if len(sys.argv) > 2:
        if sys.argv[1] == 'aug':
            waldos, notWaldos = get_images_data(sys.argv[2])
            data_augmentation(waldos, notWaldos, sys.argv[2])
        elif sys.argv[1] == 'noise':
            waldos = get_waldos_data(sys.argv[2])
            data_add_noise(waldos, sys.argv[2])
    else:
        print('Usage: data_augmentation.py [aug or noise] [image directory]')


import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
from scipy.optimize import root
from deep_sets import Deep_Sets

'''
Class for the Deep Sets Neural-Network Quantum Field State, building off the
architecture developed in https://arxiv.org/abs/2112.11957. In this ansatz,
each n-particle state is modelled as \varphi_n = DS1({x_i})*DS2({x_i-x_j}),
for deep sets networks DS1 and DS2. This ansatz is specially designed for the
Lieb-Liniger model (hence the subscript '_LL').
'''

class NQFS_LL():
    def __init__(self, DS_width, DS_depth_phi, DS_depth_rho, L, periodic):
        '''
        Initializes the NQFS ansatz

        Parameters:
          DS_width, DS_depth_phi, DS_depth_rho: width and depths of the underlying
            neural networks (i.e. rho and phi) of DS1 and DS2
          L: system length
          periodic: system periodicity (True if periodic, False if not)
        '''

        # Initializes deep sets DS1 and DS2
        input_dim1 = 2  # The input to DS1 is a 2d embedding
        self.DS1 = Deep_Sets(input_dim1, DS_width, DS_depth_phi, DS_depth_rho)
        input_dim2 = 1  # The input to DS2 is a 1d embedding
        self.DS2 = Deep_Sets(input_dim2, DS_width, DS_depth_phi, DS_depth_rho)

        # Register and initialize parameters of q_n
        '''q_n is modelled as a smoothed rectangular pulse; its parameters are
           included in DS1 for optimization. Explicitly,
           q_n = 1/(1+e^{-s*(n-c_1)})*1/(1+e^{s*(n-c_2)}) for 0 < c_1 < c_2 and s > 0,
           where q_n_mean = (c_1 + c_2)/2, q_n_width = (c_2 - c_1), q_n_slope = s.'''
        self.DS1.register_parameter(name='q_n_mean',
                                    param=torch.nn.Parameter(torch.tensor(3.)))
        # Parametrize the inverse softplus of the width to ensure width is positive
        self.DS1.register_parameter(name='q_n_inv_softplus_width',
                                    param=torch.nn.Parameter(torch.tensor(6.)))
        # Parameterize the inverse softplus of the slope to ensure slope is positive
        self.DS1.register_parameter(name='q_n_inv_softplus_slope',
                                    param=torch.nn.Parameter(torch.tensor(1.)))

        # Set system parameters
        self.L = L
        self.periodic = periodic

    def Embedding1(self, x):
        '''
        Constructs the feature embedding of particle positions x for DS1

        Parameters:
          x: an tensor of dimension (n_samples, n_particles) representing n_samples
            samples of n_particles particle positions

        Returns:
          embedded_data: a tensor of dimension (n_samples, n_particles, 2)
            containing the embedding of the particle positions
        '''

        # If periodic, embedding = (sin(2*pi*x/L), cos(2*pi*x/L))
        # If not periodic, embedding = (x/L, 1-x/L)
        x_norm = x[:, :, None] / self.L
        if self.periodic:
            sines = torch.sin(2 * np.pi * x_norm)
            cosines = torch.cos(2 * np.pi * x_norm)
            embedded_data = torch.cat((sines, cosines), axis=2)
        else:
            embedded_data = torch.cat((x_norm, 1 - x_norm), axis=2)

        return embedded_data

    def Embedding2(self, x):
        '''
        Constructs the embedding of particle positions x for DS2

        Parameters:
          x: an tensor of dimension (n_samples, n_particles) representing n_samples
            samples of n_particles particle positions

        Returns:
          embedded_data: a tensor of dimension
            (n_samples, n_particles*(n_particles-1)/2, 1) containing the embedding
            of the particle positions
        '''

        # Determines interparticle seperations {x_i - x_j} for i < j
        idx = torch.triu_indices(x.shape[1], x.shape[1], offset=1)
        interparticle_seps = (x[:, :, None] - x[:, None, :])[:, idx[0], idx[1]]

        # If periodic, embedding = cos(2*pi*(x_i-x_j)/L)
        # If not periodic, embedding = ((x_i-x_j)/L)^2
        if self.periodic:
            embedded_data = torch.cos(2 * np.pi / self.L * interparticle_seps[..., None])
        else:
            embedded_data = (interparticle_seps[..., None] / self.L) ** 2

        return embedded_data

    def log_Psi(self, x, m, g):
        '''
        Calculates log(Psi(x))

        Parameters:
          x: a tensor of dimension (n_samples, n_particles) representing
            n_samples samples of n_particles particle positions. x may contain
            entries corresponding to non-existent particles; we denote these by a
            'nan' value
          m: particle mass
          g: Lieb-Liniger interaction strength

        Returns:
          val: a tensor of dimension (n_samples, 1) containing the values of
            log_Psi(x)
        '''

        # Contribution from DS1
        x_emb1 = self.Embedding1(x)
        # Constructs mask; non-existent particle entries in x are denoted by 'nan'
        mask1 = ((~x.isnan())[..., None]).detach()
        val = self.log_DS1(x_emb1.nan_to_num(), mask1)

        # Contribution from DS2
        n = torch.sum(~x.isnan(), axis=1)[:, None]
        x_emb2 = self.Embedding2(x)
        mask2 = (~x_emb2.isnan()).detach()
        val += (n >= 2) * self.log_DS2(x_emb2.nan_to_num(), mask2)

        # Contribution from L^(-n/2)
        val += -1 / 2 * (n * np.log(self.L))

        # Contribution from Jastrow factor
        val += (n >= 2) * self.log_jastrow(x, m, g, n, mask1)

        # Contribution from q_n
        val += 1 / 2 * self.log_q_n(n)

        # Contribution from cutoff factor if system is not periodic
        if not self.periodic:
            val += self.log_cutoff_factor(x, n)

        return val

    def log_DS1(self, x_emb1, mask1):
        '''
        Returns log(DS1(x_emb1, mask1))

        Parameters:
          x_emb1: a tensor of dimension (n_samples, n_particles, 2)
            containing the embedding of the particle positions for DS1
          mask1: a Boolean tensor of dimension (n_samples, n_particles),
            denoting existent particle positions

        Returns:
          val: a tensor of dimension (n_samples, 1) containing the values of
            log(DS1(x_emb1, mask1))
        '''

        val = torch.log(self.DS1.forward(x_emb1, mask1))
        return val

    def log_DS2(self, x_emb2, mask2):
        '''
        Returns log(DS2(x_emb2, mask2))

        Parameters:
          x_emb2: a tensor of dimension (n_samples, n_particles, 1)
            containing the embedding of the particle positions for DS2
          mask2: a Boolean tensor of dimension
            (n_samples, n_particles*(n_particles-1)/2) denoting existent particle
            positions

        Returns:
          val: a tensor of dimension (n_samples, 1) containing the values of
            log(DS2(x_emb2, mask2))
        '''

        val = torch.log(self.DS2.forward(x_emb2, mask2))
        return val

    def log_jastrow(self, x, m, g, n, mask1):
        '''
        Returns log(Jastrow factor) = log(\prod_{i<j}(|x_i-x_j| + 1/(m*g)))

        Parameters:
          x: a tensor of dimension (n_samples, n_particles) containing particle
            positions
          m: particle mass
          g: Lieb-Liniger interaction strength

        Returns:
          val: a tensor of dimension (n_samples, 1) containing the values of
            log(Jastrow factor)
        '''

        # Determines interparticle seperations {x_i - x_j} for i < j
        idx = torch.triu_indices(x.shape[1], x.shape[1], offset=1)
        interparticle_seps = (x[:, :, None] - x[:, None, :])[:, idx[0], idx[1]]

        # Calculates Jastrow factor
        jastrows = torch.abs(interparticle_seps) + 1 / (m * g)
        val = torch.sum(torch.log(jastrows).nan_to_num(), axis=1)[:, None]

        # Accounts for normalization factor in Tonks-Girardeau limit
        if x.shape[1] > 0 and g > 1e4:
            log_Selberg_vals = [(2 * torch.special.gammaln(0 * n + 1 + j) + \
                                 torch.special.gammaln(0 * n + 2 + j) - \
                                 torch.special.gammaln(1 + n + j)) \
                                for j in range(0, x.shape[1])]
            log_Selberg_vals = torch.stack(log_Selberg_vals, axis=1)
            log_Selberg = torch.sum(log_Selberg_vals * mask1, axis=1)
            val += -1 / 2 * (log_Selberg)

        return val

    def log_q_n(self, n):
        '''
        Calculates log(q_n) where q_n = 1/(1+e^{-s*(n-c_1)})*1/(1+e^{s*(n-c_2)})
        for 0 < c_1 < c_2 and s > 0, and q_n_mean = (c_1 + c_2)/2,
        q_n_width = (c_2 - c_1), and  q_n_slope = s

        Parameters:
          n: a tensor of dimension (n_samples, 1) containing the number of
            (existent) particles for each sample

        Returns:
          val: a tensor of dimension (n_samples, 1) containing the values of
            log(q_n)
        '''

        q_n_width = torch.log(1 + torch.exp(self.DS1.q_n_inv_softplus_width))
        c_1 = 1 / 2 * (2 * self.DS1.q_n_mean - q_n_width)
        c_2 = 1 / 2 * (2 * self.DS1.q_n_mean + q_n_width)
        q_n_slope = torch.log(1 + torch.exp(self.DS1.q_n_inv_softplus_slope))
        s = q_n_slope

        val = -torch.log(1 + torch.exp(-s * (n - c_1)))
        val += -torch.log(1 + torch.exp(s * (n - c_2)))

        return val

    def log_cutoff_factor(self, x, n):
        '''
        Calculates logarithm of the cutoff factor used to implement closed boundary
        conditions. Explicitly, cutoff factor = 1/N*\prod_i{x_i/L*(1-x_i/L)},
        where 1/N = (30/L)^(n/2) is a normalization factor.

        Parameters:
          x: a tensor of dimension (n_samples, n_particles) representing
            n_samples samples of n_particles particle positions
          n: a tensor of dimension (n_samples, 1) containing the number of
            (existent) particles of each sample

        Returns:
          val: a tensor of dimension (n_samples, 1) that contains the values of
            log(cutoff factor)
        '''

        val = (x % self.L) / self.L * (1 - (x % self.L) / self.L)
        val = torch.sum(torch.log(val + 1e-16).nan_to_num(), axis=1)[:, None]
        val += -n * 1 / 2 * np.log(self.L / 30)

        return val

    def MetropolisHastings_FockSpace(self, n_samples, n_chains, p_pm, n_0,
                                     GPU_device, m, g):
        '''
        Performs the Metropolis-Hastings algorithm in Fock space to produce
        samples of particle configurations drawn from P_n and |\varphi_n|^2. This
        runs multiple MCMC chains on GPU to parallelize data generation.

        Parameters:
          n_samples: number of samples drawn for each MCMC chain
          n_chains: number of separate MCMC chains
          p_pm: probability of increasing or decreasing particle number at each
            configuration proposal
          n_0: initial particle number of each MCMC chain
          GPU_device: the GPU device under use (i.e. cuda)

        Returns:
          x_list: a list of n_samples tensors, each containing the sampled particle
            configurations across the MCMC chains at the specified iteration. Each
            such tensor is of dimension (n_chains, n), where n is the maximum
            particle number of the sampled configurations at the corresponding
            iteration.
          n_list: a list of n_samples tensors, each containing the number of
            particles of the sampled particle configurations across the MCMC
            chains at the specified iteration. Each such tensor is of dimension
            (n_chains, 1), and takes integer values.
        '''

        # Sets width of the uniform distribution used for configuration proposals
        w = 0.3 * self.L

        # Initializes n_list as the number of particles, each configuration
        # containing n_0 particles at the beginning
        n_list = [n_0 * torch.ones(n_chains, dtype=torch.long, device=GPU_device)]

        # Initializes x_list as particle positions; positions are chosen randomly; a
        # value 'nan' is used to denote the absence of a particle
        if n_0 == 0:
            x_0 = torch.cuda.FloatTensor(n_chains, 1).fill_(float('nan'))
        else:
            x_0 = self.L * torch.rand(n_chains, n_0, device=GPU_device)
        x_list = [x_0]

        # Sets current particle positions, particle number, and log(Psi)
        x = x_list[0].detach().clone()
        n = n_list[0].detach().clone()
        log_Psi_val = self.log_Psi(x, m, g)

        # Runs MCMC
        for i in range(n_samples - 1):
            # Generates random numbers to choose which chains will be proposed to have
            # a particle added or removed; if neither, the proposal perturbs the
            # particle positions
            u = torch.cuda.FloatTensor(2, n_chains).uniform_()
            add_ind = torch.where(u[0, :] < p_pm)[0]
            remove_ind = torch.where(torch.logical_and(u[0, :] > p_pm,
                                                       u[0, :] < 2 * p_pm))[0]
            perturb_ind = (u[0, :] > 2 * p_pm)[:, None]

            # Generates proposed configurations

            # Adds a particle at a randomly chosen position
            x_proposed = torch.cat((x, \
                                    torch.cuda.FloatTensor(n_chains, 1).fill_(float('nan'))), axis=1)
            x_proposed[add_ind, n[add_ind]] = \
                self.L * torch.cuda.FloatTensor(add_ind.shape[0]).uniform_()

            # Removes a particle at random (n_cutoff is used to ensure particle number
            # doesn't become negative)
            n_cutoff = torch.maximum(torch.cuda.FloatTensor(1).fill_(1.),
                                     n[remove_ind]).long()
            x_remove_ind = torch.floor( \
                torch.cuda.FloatTensor(remove_ind.shape[0]).uniform_() * n_cutoff).long()
            x_proposed[remove_ind, x_remove_ind] = x_proposed[remove_ind, n_cutoff - 1]
            x_proposed[remove_ind, n_cutoff - 1] = float('nan')

            # Perturbs particle positions by a uniform random variable of mean 0 and
            # width w
            perturbation = torch.cuda.FloatTensor(x_proposed.shape).uniform_(-w / 2, w / 2)
            x_proposed += perturb_ind * perturbation
            x_proposed %= self.L

            # Calculates particle numbers of proposed configurations
            n_proposed = n.detach().clone()
            n_proposed[add_ind] += 1
            n_proposed[remove_ind] = torch.maximum(n_proposed[remove_ind] - 1, \
                                                   torch.cuda.FloatTensor(1).fill_(0.).long())

            # Computes acceptance ratio
            log_Psi_val_proposed = self.log_Psi(x_proposed, m, g)
            L_factor = torch.cuda.FloatTensor(n_chains, 1).fill_(1.)
            L_factor[add_ind, 0] *= self.L
            L_factor[remove_ind, 0] *= 1 / self.L
            A = torch.minimum(torch.cuda.FloatTensor(n_chains, 1).fill_(1.),
                              L_factor * torch.exp(2 * (log_Psi_val_proposed - log_Psi_val)))

            # Accepts or rejects configuration proposals (using the random numbers u
            # generated earlier), and then updates x, n, and log_Psi_val
            accept_ind = torch.where(u[1, :] < A[:, 0])[0]
            x = torch.cat((x, \
                           torch.cuda.FloatTensor(n_chains, 1).fill_(float('nan'))), axis=1)
            x[accept_ind, :] = x_proposed[accept_ind, :]
            while torch.all(x[:, -1].isnan()) and x.shape[1] > 1:
                x = x[:, :-1]
            log_Psi_val[accept_ind, :] = log_Psi_val_proposed[accept_ind, :]
            n[accept_ind] = n_proposed[accept_ind]

            # Records particle configurations and particle number
            n_list.append(n.detach().clone())
            x_list.append(x.detach().clone())

        return x_list, n_list

    def GenerateMCMCSamples_FockSpace(self, n_samples, n_chains, p_pm, n_0,
                                      GPU_device, m, g):
        '''
        Generates particle configurations drawn from P_n and |\varphi_n|^2 using the
        algorithm for Metropolis-Hastings in Fock space, and subsequently removes
        MCMC burn-in phase and sorts the configurations by particle number.

        Parameters:
          n_samples: number of samples drawn for each MCMC chain
          n_chains: number of separate MCMC chains
          p_pm: probability of increasing or decreasing particle number at each
            configuration proposal
          n_0: initial particle number of each MCMC chain
          GPU_device: the GPU device under use (i.e. cuda)

        Returns:
          x_sorted: a list of tensors, the nth of which contains the sampled
            particle configurations containing n particles and is of dimension
            (#, n), where # denotes the number of drawn samples of particle number n
          n: a tensor containing the particle numbers of the configurations drawn
            from P_n and |\varphi_n|^2
          chain_idx_sorted: a list of tensors, each of the same dimension as those
            of x_sorted. An entry of such tensor denotes the MCMC chain from
            which the corresponding particle configuration came.
        '''

        # Generate samples
        x_list, n_list = self.MetropolisHastings_FockSpace(n_samples, n_chains, p_pm,
                                                           n_0, GPU_device, m, g)

        # Discard values from the MCMC burn-in phase
        discard_fraction = 0.3
        n_discard = min(500, int(discard_fraction * n_samples))
        n = torch.flatten(torch.stack(n_list[n_discard:]))
        n_max = torch.max(n)
        n_min = torch.min(n)
        x = torch.cuda.FloatTensor((n_samples - n_discard) * n_chains,
                                   max(1, n_max)).fill_(float('nan'))
        for i in range(n_samples - n_discard):
            x[i * n_chains:(i + 1) * n_chains, :x_list[i + n_discard].shape[1]] = \
                x_list[i + n_discard]

            # Sort by particle number and determine chain of origin
        x_sorted = []
        chain_idx = torch.arange(n_chains, \
                                 device=GPU_device)[:, None].repeat(n_samples - n_discard, 1)
        chain_idx_sorted = []
        for n_val in range(n_min, n_max + 1):
            idx = torch.where(n == n_val)[0]
            x_sorted.append(x[idx, :n_val])
            chain_idx_sorted.append(chain_idx[idx, 0])

        return x_sorted, n, chain_idx_sorted

    def Local_Energy(self, x, m, mu, V, g):
        '''
        Calculates the local energies of the configurations in x

        Parameters:
          x: a tensor of dimension (n_samples, n_particles) representing
            n_samples samples of n_particles particle positions
          m: particle mass
          mu: chemical potential
          V: External potential function
          g: Lieb-Liniger interaction strength

        Returns:
          E_local: a tensor of dimension (n_samples, 1) containing the local
            energies of the configurations in x
        '''

        # Computes local energies
        if torch.numel(x) > 0:
            # Creates temporary variable to compute gradients
            x_tmp = x.detach().clone()
            x_tmp.requires_grad = True
            x_tmp.grad = None

            # Compute gradient
            tmp = self.log_Psi(x_tmp, m, g)
            tmp2 = torch.autograd.grad(outputs=tmp, inputs=x_tmp, \
                                       grad_outputs=torch.ones_like(tmp), create_graph=True)[0]
            gradient = tmp2.detach()
            # Compute second derivatives
            tmp3 = [torch.autograd.grad(outputs=tmp2[:, i], inputs=x_tmp,
                                        grad_outputs=torch.ones_like(tmp2[:, i]),
                                        retain_graph=True)[0][:, i].detach()
                    for i in range(x_tmp.shape[1])]
            second_ders = torch.stack(tmp3, axis=1)
            del tmp, tmp2, tmp3, x_tmp

            KE_local = 1 / (2 * m) * torch.sum(-gradient ** 2 - second_ders, axis=1)[:, None]
            del gradient, second_ders

            PE_local = torch.sum(V(x) - mu, axis=1)[:, None]
        else:
            KE_local = torch.cuda.FloatTensor(x.shape[0], 1).fill_(0.)
            PE_local = torch.cuda.FloatTensor(x.shape[0], 1).fill_(0.)

        E_local = KE_local + PE_local
        del KE_local, PE_local
        return E_local

    def Energy_Estimate(self, x_sorted, chain_idx_sorted, n_chains, GPU_device,
                        m, mu, V, g):
        '''
        Estimates the energies of the configurations in x_sorted

        Parameters:
          x_sorted: a list of tensors, the nth of which contains the sampled
            particle configurations of n particles and is of dimension (#, n),
            where # denotes the number of samples drawn of particle number n
          chain_idx_sorted: a list of tensors denoting the MCMC chain from
            which the corresponding particle configuration in x_sorted came.
          n_chains: number of MCMC chains
          GPU_device: the GPU device under use (i.e. cuda)
          m: particle mass
          mu: chemical potential
          V: External potential function
          g: Lieb-Liniger interaction strength

        Returns:
          E: The estimated energy
          E_std: The estimated standard deviation of E
        '''

        # Calculates local energies and the corresponding estimated energy
        Local_Energies = []
        for x in x_sorted:
            Local_Energies.append(self.Local_Energy(x, m, mu, V, g))
        Local_Energies = torch.cat(Local_Energies)
        E = torch.mean(Local_Energies)

        # Calculates standard deviation of E by binning across the MCMC chains
        chain_idx_sorted_flat = torch.cat(chain_idx_sorted)
        E_means_chains = \
            torch.tensor([torch.sum(Local_Energies[:, 0] * (chain_idx_sorted_flat == i)) / \
                          torch.sum(chain_idx_sorted_flat == i) \
                          for i in range(n_chains)], device=GPU_device)
        E_std = torch.std(E_means_chains) / (n_chains ** 0.5)

        return E, E_std

    def Local_Gradient_Energy(self, x, m, mu, V, g):
        '''
        Calculates the local energies and the local gradient of the energy w.r.t.
        neural network parameters

        Parameters:
          x: a tensor of dimension (n_samples, n_particles) representing
            n_samples samples of n_particles particle positions
          m: particle mass
          mu: chemical potential
          V: External potential function
          g: Lieb-Liniger interaction strength

        Returns:
          E_local: a tensor of dimension (n_samples, 1) containing the local
            energies of the configurations in x
          log_Psi_grad_sum1, log_Psi_grad_sum2 : a tensor of dimension
            (n_samples, 1) containing the sum of the gradient of the logarithm of
            Psi w.r.t. the parameters of DS1, DS2, respectively, for configurations
            in x
          E_log_Psi_grad_sum1, E_log_Psi_der_sum2: a tensor of dimension
            (n_samples, 1) containing the sum of the local energy times the gradient
            of the logarithm of Psi w.r.t. the parameters of DS1, DS2, respectively,
            for the configurations in x
        '''

        # Computes local energies
        if torch.numel(x) > 0:
            # Creates temporary variable to compute gradients
            x_tmp = x.detach().clone()
            x_tmp.requires_grad = True
            x_tmp.grad = None

            # Compute gradient
            tmp = self.log_Psi(x_tmp, m, g)
            tmp2 = torch.autograd.grad(outputs=tmp, inputs=x_tmp, \
                                       grad_outputs=torch.ones_like(tmp), create_graph=True)[0]
            gradient = tmp2.detach()
            # Compute second derivatives
            tmp3 = [torch.autograd.grad(outputs=tmp2[:, i], inputs=x_tmp,
                                        grad_outputs=torch.ones_like(tmp2[:, i]),
                                        retain_graph=True)[0][:, i].detach()
                    for i in range(x_tmp.shape[1])]
            second_ders = torch.stack(tmp3, axis=1)
            del tmp, tmp2, tmp3, x_tmp

            KE_local = 1 / (2 * m) * torch.sum(-gradient ** 2 - second_ders, axis=1)[:, None]
            del gradient, second_ders

            PE_local = torch.sum(V(x) - mu, axis=1)[:, None]
        else:
            KE_local = torch.cuda.FloatTensor(x.shape[0], 1).fill_(0.)
            PE_local = torch.cuda.FloatTensor(x.shape[0], 1).fill_(0.)

        E_local = KE_local + PE_local
        del KE_local, PE_local

        # Computes the sum of the gradient of log_Psi wrt neural network parameters
        x.grad = None
        tmp = self.log_Psi(x, m, g)
        # Accumulated gradient is naturally the sum
        log_Psi_grad_sum1 = torch.autograd.grad(outputs=tmp,
                                                inputs=self.DS1.parameters(),
                                                grad_outputs=torch.ones_like(tmp))
        tmp = self.log_Psi(x, m, g)
        log_Psi_grad_sum2 = torch.autograd.grad(outputs=tmp,
                                                inputs=self.DS2.parameters(),
                                                grad_outputs=torch.ones_like(tmp))

        # Comptues the sum of the local energy times the gradient of the logarithm
        # of Psi w.r.t. neural network parameters
        x.grad = None
        tmp = self.log_Psi(x, m, g)
        tmp = E_local * tmp
        E_log_Psi_grad_sum1 = torch.autograd.grad(outputs=tmp,
                                                  inputs=self.DS1.parameters(),
                                                  grad_outputs=torch.ones_like(tmp))
        x.grad = None
        tmp = self.log_Psi(x, m, g)
        tmp = E_local * tmp
        E_log_Psi_grad_sum2 = torch.autograd.grad(outputs=tmp,
                                                  inputs=self.DS2.parameters(),
                                                  grad_outputs=torch.ones_like(tmp))

        return E_local, log_Psi_grad_sum1, log_Psi_grad_sum2, E_log_Psi_grad_sum1, E_log_Psi_grad_sum2

    def Energy_Grad_Estimate(self, x_sorted, n, chain_idx_sorted, n_chains,
                             GPU_device, m, mu, V, g):
        '''
        Estimates the gradient of the energy w.r.t. neural network parameters. Also
        estimates the energy and its standard deviation.

        Parameters:
          x_sorted: a list of tensors, the nth of which contains the sampled
            particle configurations containing n particles and is of dimension
            (#, n), where # denotes the number of samples drawn of particle number n
          n: a tensor containing the particle numbers of configurations in x_sorted
          chain_idx_sorted: a list of tensors denoting the MCMC chain from
            which the corresponding particle configuration in x_sorted came.
          n_chains: number of chains
          GPU_device: the GPU device under use (i.e. cuda)
          m: particle mass
          mu: chemical potential
          V: external potential function
          g: Lieb-Liniger interaction strength

        Returns:
          E: the estimated energy
          E_std: the estimated standard deviation of E
          E_grad1, E_grad2: gradient of E wrt parameters of DS1, DS2, respectively
        '''

        # Initializes variables
        Local_Energies = []
        log_Psi_grad1 = []
        log_Psi_grad2 = []
        E_log_Psi_grad1 = []
        E_log_Psi_grad2 = []
        for param in self.DS1.parameters():
            log_Psi_grad1.append(torch.cuda.FloatTensor(param.shape).fill_(0.))
            E_log_Psi_grad1.append(torch.cuda.FloatTensor(param.shape).fill_(0.))
        for param in self.DS2.parameters():
            log_Psi_grad2.append(torch.cuda.FloatTensor(param.shape).fill_(0.))
            E_log_Psi_grad2.append(torch.cuda.FloatTensor(param.shape).fill_(0.))

            # Calculates local energies and local energy gradients
        for x in x_sorted:
            Local_Energies_tmp, log_Psi_grad_sum1, log_Psi_grad_sum2, \
            E_log_Psi_grad_sum1, E_log_Psi_grad_sum2 = \
                self.Local_Gradient_Energy(x, m, mu, V, g)
            Local_Energies.append(Local_Energies_tmp)
            for i in range(len(list(self.DS1.parameters()))):
                log_Psi_grad1[i] += log_Psi_grad_sum1[i]
                E_log_Psi_grad1[i] += E_log_Psi_grad_sum1[i]
            for i in range(len(list(self.DS2.parameters()))):
                log_Psi_grad1[i] += log_Psi_grad_sum1[i]
                E_log_Psi_grad1[i] += E_log_Psi_grad_sum1[i]

        for i in range(len(list(self.DS1.parameters()))):
            log_Psi_grad1[i] *= 1 / n.shape[0]
            E_log_Psi_grad1[i] *= 1 / n.shape[0]
        for i in range(len(list(self.DS2.parameters()))):
            log_Psi_grad2[i] *= 1 / n.shape[0]
            E_log_Psi_grad2[i] *= 1 / n.shape[0]

            # Estimates energy and its standard deviations
        Local_Energies = torch.cat(Local_Energies)
        E = torch.mean(Local_Energies)
        chain_idx_sorted_flat = torch.cat(chain_idx_sorted)
        E_means_chains = \
            torch.tensor([torch.sum(Local_Energies[:, 0] * (chain_idx_sorted_flat == i)) / \
                          torch.sum(chain_idx_sorted_flat == i) \
                          for i in range(n_chains)], device=GPU_device)
        E_std = torch.std(E_means_chains) / (n_chains ** 0.5)

        # Estimates gradients of energy wrt neural network parameters
        E_grad1 = []
        E_grad2 = []
        for i in range(len(list(self.DS1.parameters()))):
            tmp = 2 * (E_log_Psi_grad1[i] - E * log_Psi_grad1[i])
            E_grad1.append(tmp)
        for i in range(len(list(self.DS2.parameters()))):
            tmp = 2 * (E_log_Psi_grad2[i] - E * log_Psi_grad2[i])
            E_grad2.append(tmp)

        return E, E_std, E_grad1, E_grad2

    def ADAM_update(self, x_sorted, n, chain_idx_sorted, n_chains, GPU_device, m, mu, V, g, t, m1, v1,
                    m2, v2, beta_1, beta_2, lr, lr_q):
        '''
        Performs a single ADAM update

        Parameters:
          x_sorted: a list of tensors, the nth of which contains the sampled
            particle configurations containing n particles and is of dimension
            (#, n), where # denotes the number of samples drawn of particle number n
          n: a tensor containing the particle numbers of configurations in x_sorted
          chain_idx_sorted: a list of tensors denoting the MCMC chain from
            which the corresponding particle configuration in x_sorted came.
          chain_idx_sorted: a list of tensors, each of the same dimension as those
            of x_sorted. An entry of such tensor denotes the MCMC chain from
            which the corresponding particle configuration came.
          n_chains: number of separate MCMC chains
          GPU_device: the GPU device under use (i.e. cuda)
          m: particle mass
          mu: chemical potential
          V: external potential function
          g: Lieb-Liniger interaction strength
          t: current ADAM iteration
          m1, m2: first moment estimate for DS1, DS2, respectively
          v1, v2: second moment estimate for DS1, DS2, respectively
          beta_1, beta_2: ADAM parameters
          lr: learning rate for parameters of DS1, DS
          lr_q: learning rate for parameters of q_n

        Returns:
          E: the estimated energy
          E_std: the estimated standard deviation of E
          m1, m2: updated first moment estimate for DS1, DS2, respectively
          v1, v2: updated second moment estimate for DS1, DS2, respectively
        '''

        eps = 1e-8  # ADAM parameter

        # Estimates gradients
        E, E_std, E_grad1, E_grad2 = self.Energy_Grad_Estimate(x_sorted, n, \
                                                               chain_idx_sorted, n_chains, GPU_device, m, mu, V, g)

        # Performs ADAM update on DS1
        for i, param in enumerate(self.DS1.parameters()):
            m1[i] = beta_1 * m1[i] + (1 - beta_1) * E_grad1[i]
            v1[i] = beta_2 * v1[i] + (1 - beta_2) * E_grad1[i] ** 2
            m1_hat = m1[i] / (1 - beta_1 ** t)
            v1_hat = v1[i] / (1 - beta_2 ** t)
            if i < 3:
                param.data -= lr_q * m1_hat / (torch.sqrt(v1_hat) + eps)
            else:
                param.data -= lr * m1_hat / (torch.sqrt(v1_hat) + eps)

        # Performs ADAM update on DS2
        for i, param in enumerate(self.DS2.parameters()):
            m2[i] = beta_1 * m2[i] + (1 - beta_1) * E_grad2[i]
            v2[i] = beta_2 * v2[i] + (1 - beta_2) * E_grad2[i] ** 2
            m2_hat = m2[i] / (1 - beta_1 ** t)
            v2_hat = v2[i] / (1 - beta_2 ** t)
            param.data -= lr * m2_hat / (torch.sqrt(v2_hat) + eps)

        return E, E_std, m1, v1, m2, v2

    def minimize_energy_ADAM(self, n_samples, n_chains, p_pm, GPU_device,
                             m, mu, V, g, beta_1, beta_2, lr, lr_q, n_iters):
        '''
        Minimizes energy with ADAM, and outputs the estimated energy and particle
        number at each iteration.

        Parameters:
          n_samples: number of samples drawn for each MCMC chain
          n_chains: number of separate MCMC chains
          p_pm: probability of increasing or decreasing particle number at each
            configuration proposal
          GPU_device: the GPU device under use (i.e. cuda)
          m: particle mass
          mu: chemical potential
          V: external potential function
          g: Lieb-Liniger interaction strength
          beta_1, beta_2: ADAM parameters
          lr: learning rate for parameters of DS1, DS
          lr_q: learning rate for parameters of q_n
          n_iters: number of iterations of ADAM to be performed

        Returns:
          Es: an array containing the estimated energies at each iteration
          E_stds: an array containing the estimated standard deviations of energy
            at each iteration
          n_means: an array containing the estimated mean particle number at each
            iteration
          n_stds: an array containing the estimated standard deviation of particle
            number at each iteration
        '''

        # Initializes variables
        Es = torch.cuda.FloatTensor(n_iters).fill_(0.)
        E_stds = torch.cuda.FloatTensor(n_iters).fill_(0.)
        n_means = torch.cuda.FloatTensor(n_iters).fill_(0.)
        n_stds = torch.cuda.FloatTensor(n_iters).fill_(0.)
        m1, v1 = [], []
        for param in self.DS1.parameters():
            m1.append(torch.cuda.FloatTensor(param.shape).fill_(0.))
            v1.append(torch.cuda.FloatTensor(param.shape).fill_(0.))
        m2, v2 = [], []
        for param in self.DS2.parameters():
            m2.append(torch.cuda.FloatTensor(param.shape).fill_(0.))
            v2.append(torch.cuda.FloatTensor(param.shape).fill_(0.))
        n_means[-1] = int(round(self.DS1.q_n_mean.data.item()))

        # Performs ADAM for n_iters number of iterations
        for t in range(1, n_iters + 1):
            # Generates MCMC data
            n_0 = int(n_means[t - 2].item())
            x_sorted, n, chain_idx_sorted = \
                self.GenerateMCMCSamples_FockSpace(n_samples, n_chains, p_pm, n_0,
                                                   GPU_device, m, g)

            # Performs ADAM update
            E, E_std, m1, v1, m2, v2 = \
                self.ADAM_update(x_sorted, n, chain_idx_sorted, n_chains, GPU_device, m, mu, V, g, t, m1, v1,
                                 m2, v2, beta_1, beta_2, lr, lr_q)

            # Stores energy, particle number, and the standard deviations thereof
            Es[t - 1] = E.detach().clone()
            E_stds[t - 1] = E_std.detach().clone()
            n_mean = torch.mean(n.to(torch.float))
            n_std = torch.std(n.to(torch.float))
            n_means[t - 1] = n_mean.detach().clone()
            n_stds[t - 1] = n_std.detach().clone()

            # Prints results of current iteration
            print("Iteration: " + str(t) + "/" + str(n_iters))
            print("Energy: " + str(round(E.item(), 3)) + " +- " \
                  + str(round(E_std.item(), 3)))
            print("Number of particles: " + str(round(n_mean.item(), 3)) + \
                  " +- " + str(round(n_std.item(), 3)))
            print("\n")

        return Es, E_stds, n_means, n_stds

    def plot_energy(self, Es, E_stds, E_exact, start, end):
        '''
        Plots energy vs. iteration from start to end.

        Parameters:
          Es: an array containing the estimated energies at each iteration
          E_stds: an array containing the estimated standard deviations of energy
          E_exact: the exact ground state energy
          start: the starting iteration from which data is plotted
          end: the last iteration from which data is plotted
        '''

        iterations = np.arange(start, end)

        # Plots energy
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(6.5, 2.8), dpi=200)
        ax1.plot(iterations, Es.cpu()[start:end], linewidth=.7, label='Energy')
        ax1.fill_between(iterations, (Es - E_stds).cpu()[start:end],
                         (Es + E_stds).cpu()[start:end], facecolor='blue',
                         alpha=0.25, interpolate=True)
        ax1.plot(iterations, E_exact + 0 * Es.cpu()[start:end], linewidth=1,
                 label='Exact Energy')
        ax1.set_title('Energy vs. Iteration', fontsize=10)
        ax1.set_ylabel('Energy', fontsize=8)
        ax1.set_xlabel('Iteration', fontsize=8)
        ax1.legend(prop={'size': 7})
        ax1.set_xticklabels(ax1.get_xticks().astype(int), fontsize=6)
        ax1.set_yticklabels(np.round(ax1.get_yticks(), 2), fontsize=6)

        # Plots error in energy and relative standard deviation
        Energy_errors = (torch.abs((Es.cpu() - E_exact) / E_exact))[start:end]
        ax2.plot(iterations, Energy_errors, linewidth=.7, label='Relative Error')
        ax2.plot(iterations, (torch.abs(E_stds / Es)).cpu()[start:end],
                 linewidth=.7, label='Relative St. Dev.')
        ax2.set_yscale('log')
        ax2.set_title('Energy Discrepancies vs. Iteration', fontsize=10)
        ax2.set_ylabel('Value', fontsize=8)
        ax2.set_xlabel('Iteration', fontsize=8)
        ax2.legend(prop={'size': 7})
        ax2.set_xticklabels(ax2.get_xticks().astype(int), fontsize=6)
        ax2.set_yticklabels(np.round(ax2.get_yticks(), 10), fontsize=6)

        fig.tight_layout()
        fig.show()

    def plot_n(self, n_means, n_stds, n_exact, start, end):
        '''
        Plots particle number vs. iteration from start to end

        Parameters:
          n_means: an array containing the estimated mean particle number at each
            iteration
          n_stds: an array containing the estimated standard deviation of particle
            number at each iteration
          n_exact: the exact ground state energy
          start: the starting iteration from which data is plotted
          end: the last iteration from which data is plotted
        '''

        iterations = np.arange(start, end)

        # Plots particle number
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(6.5, 2.8), dpi=200)
        ax1.plot(iterations, n_means.cpu()[start:end], linewidth=1, label='n')
        ax1.fill_between(iterations, (n_means - n_stds).cpu()[start:end],
                         (n_means + n_stds).cpu()[start:end], facecolor='blue',
                         alpha=0.25, interpolate=True)
        ax1.plot(iterations, n_exact + 0 * n_means.cpu()[start:end],
                 linewidth=1, label='Exact n')
        ax1.set_title('Particle Number vs. Iteration', fontsize=10)
        ax1.set_ylabel('n', fontsize=8)
        ax1.set_xlabel('Iteration', fontsize=8)
        ax1.legend(prop={'size': 7})
        ax1.set_xticklabels(ax1.get_xticks().astype(int), fontsize=6)
        ax1.set_yticklabels(np.round(ax1.get_yticks(), 2), fontsize=6)

        # Plots error in particle number and relative standard deviation
        n_errors = (torch.abs((n_means.cpu() - n_exact) / n_exact))[start:end]
        ax2.plot(iterations, n_errors, linewidth=.7, label='Relative Error')
        ax2.plot(iterations, (n_stds / n_means).cpu()[start:end],
                 linewidth=.7, label='Relative St. Dev.')
        ax2.set_yscale('log')
        ax2.set_title('Particle Number Discrepancies vs. Iteration', fontsize=10)
        ax2.set_ylabel('Value', fontsize=8)
        ax2.set_xlabel('Iteration', fontsize=8)
        ax2.legend(prop={'size': 7})
        ax2.set_xticklabels(ax2.get_xticks().astype(int), fontsize=6)
        ax2.set_yticklabels(np.round(ax2.get_yticks(), 10), fontsize=6)

        fig.tight_layout()
        fig.show()

    def Number_KE_density_function(self, N_pts, x_sorted, chain_idx_sorted, n_chains,
                                   GPU_device, m, g, N_int, batch_size):
        '''
        Estimates the number density function and kinetic energy density function
        across the range [0,L], and the standard deviations thereof

        Parameters:
          N_pts: number of points at which the number density is estimated over
            the range [0,L]
          x_sorted: a list of tensors, the nth of which contains the sampled
            particle configurations of n particles and is of dimension (#, n),
            where # denotes the number of samples drawn of particle number n
          chain_idx_sorted: a list of tensors denoting the MCMC chain from
            which the corresponding particle configuration in x_sorted came.
          n_chains: number of MCMC chains
          GPU_device: the GPU device under use (i.e. cuda)
          N_int: number of integration points used to compute the numerical integral
            needed for local number density calculation
          batch_size: the size of batches over which the local number density is
            estimated; performing the numerical integral is memory intensive, and
            hence we batch over entries to save memory

        Returns:
          num_densities: the estimated number densities
          num_densities_std: the standard deviation of the estimated
            number densities
          KE_densities: the estimated kinetic energy densities
          KE_densities_std: the standard deviation of the estimated kinetic
            energy densities
        '''

        # Evaluates number density and kinetic energy density across N_pts points
        xs = np.linspace(0, self.L, N_pts)
        num_densities = np.zeros(N_pts)
        num_densities_std = np.zeros(N_pts)
        KE_densities = np.zeros(N_pts)
        KE_densities_std = np.zeros(N_pts)
        for i in range(N_pts):
            print('Evaluating number density and KE density at point ' + str(i + 1) + \
                  '/' + str(N_pts))
            x_eval = xs[i]
            num_density, num_density_std, KE_density, KE_density_std = \
                self.Number_KE_density(x_eval, x_sorted, chain_idx_sorted, n_chains, \
                                       GPU_device, m, g, N_int, batch_size)
            num_densities[i] = num_density
            num_densities_std[i] = num_density_std
            KE_densities[i] = KE_density
            KE_densities_std[i] = KE_density_std

        return num_densities, num_densities_std, KE_densities, KE_densities_std

    def Number_KE_density(self, x_eval, x_sorted, chain_idx_sorted, n_chains,
                          GPU_device, m, g, N_int, batch_size):
        '''
        Estimates the number density and kinetic energy density at x_eval and their
        standard deviations

        Parameters:
          x_eval: a scalar value of the point at which the number density is to be
            estimated
          x_sorted: a list of tensors, the nth of which contains the sampled
            particle configurations of n particles and is of dimension (#, n),
            where # denotes the number of samples drawn of particle number n
          chain_idx_sorted: a list of tensors denoting the MCMC chain from
            which the corresponding particle configuration in x_sorted came.
          n_chains: number of MCMC chains
          GPU_device: the GPU device under use (i.e. cuda)
          N_int: number of integration points used to compute the numerical integral
            needed for local number density calculation
          batch_size: the size of batches over which the local number density is
            estimated; performing the numerical integral is memory intensive, and
            hence we batch over entries to save memory

        Returns:
          num_density: the estimated number density
          num_density_std: the standard deviation of the estimated number density
          KE_density: the estimated kinetic energy density
          KE_density_std: the standard deviation of the estimated kinetic
            energy density
        '''

        # Calculates local number densities and local KE densities
        Local_num_densities = []
        Local_KE_densities = []
        for x in x_sorted:
            num_density_local, KE_density_local = self.Local_number_KE_density(x_eval, \
                                                                               x, GPU_device, m, g, N_int, batch_size)
            Local_num_densities.append(num_density_local)
            Local_KE_densities.append(KE_density_local)

        # Estimated number density and KE density
        Local_num_densities = torch.cat(Local_num_densities)
        num_density = torch.mean(Local_num_densities)
        Local_KE_densities = torch.cat(Local_KE_densities)
        KE_density = torch.mean(Local_KE_densities)

        # Calculates standard deviations by binning across the MCMC chains
        chain_idx_sorted_flat = torch.cat(chain_idx_sorted)
        num_density_means_chains = \
            torch.tensor([torch.sum(Local_num_densities[:, 0] * (chain_idx_sorted_flat == i)) / \
                          torch.sum(chain_idx_sorted_flat == i) \
                          for i in range(n_chains)], device=GPU_device)
        num_density_std = torch.std(num_density_means_chains) / (n_chains ** 0.5)
        KE_density_means_chains = \
            torch.tensor([torch.sum(Local_KE_densities[:, 0] * (chain_idx_sorted_flat == i)) / \
                          torch.sum(chain_idx_sorted_flat == i) \
                          for i in range(n_chains)], device=GPU_device)
        KE_density_std = torch.std(KE_density_means_chains) / (n_chains ** 0.5)

        return num_density, num_density_std, KE_density, KE_density_std

    def Local_number_KE_density(self, x_eval, x, GPU_device, m, g,
                                N_int, batch_size):
        '''
        Calculates the local number densities and kinetic energy densities of the
        configurations in x

        Parameters:
          x_eval: a scalar value of the point at which the number density is to be
            estimated
          x: a tensor of dimension (n_samples, n_particles) representing
            n_samples samples of n_particles particle positions
          GPU_device: the GPU device under use (i.e. cuda)
          m: particle mass
          g: Lieb-Liniger interaction strength
          N_int: number of integration points used to compute the numerical integral
            \int dx_i |\varphi_n|^2 and evaluate the marginal distribution
          batch_size: the size of batches over which the local number density is
            estimated; performing the numerical integral is memory intensive, and
            hence we batch over entries to save memory

        Returns:
          num_density_local: a tensor of dimension (n_samples, 1) containing the
            local number densities at x_eval of the configurations in x
          KE_density_local: a tensor of dimension (n_samples, 1) containing the
            local kinetic energy densities at x_eval of the configurations in x
        '''

        x_eval_tensor = torch.cuda.FloatTensor(1).fill_(x_eval)
        # The points across which the numerical integral is calculated
        x_int_points = torch.linspace(0, self.L, N_int, device=GPU_device)

        num_density_local = []
        KE_density_local = []
        for i in range(np.ceil(x.shape[0] / batch_size).astype(int)):
            # Computes log_Psi for numerator
            x_tmp = x[i * batch_size:(i + 1) * batch_size, :]
            x_tmp_numerator = torch.cat((x_eval_tensor.repeat(x_tmp.shape[0], 1),
                                         x_tmp[:, 1:]), axis=1)
            log_Psi_numerator_val = self.log_Psi(x_tmp_numerator, m, g)[:, 0].detach()

            x_tmp_numerator2 = x_tmp_numerator.detach().clone()
            x_tmp_numerator2.requires_grad = True
            x_tmp_numerator2.grad = None

            tmp = torch.exp(self.log_Psi(x_tmp_numerator2, m, g))
            tmp2 = torch.autograd.grad(outputs=tmp, inputs=x_tmp_numerator2,
                                       grad_outputs=torch.ones_like(tmp),
                                       create_graph=True)[0]
            gradient = tmp2.detach()
            del tmp, tmp2

            # Computes log_Psi for integral over distribution
            x_int = torch.cat((x_int_points[:, None, None].repeat(1, x_tmp.shape[0], 1),
                               x_tmp[None, :, 1:].repeat(N_int, 1, 1)), axis=2)
            x_int = x_int.reshape(N_int * x_tmp.shape[0], -1)
            log_Psi_integrand_vals = self.log_Psi(x_int, m, g).detach()
            log_Psi_integrand_vals = log_Psi_integrand_vals.reshape(N_int,
                                                                    x_tmp.shape[0], 1)

            # Calculates local number density
            integrand_vals = torch.exp(2 * log_Psi_integrand_vals[:, :, 0])
            integrals = torch.trapz(integrand_vals, dx=self.L / (N_int - 1), dim=0)
            num_density_local.append(x.shape[1] * torch.exp(2 * log_Psi_numerator_val) \
                                     / integrals)
            # KE_density_local.append(-1/(2*m)*x.shape[1]*torch.exp(log_Psi_numerator_val)*\
            # second_ders[:,0]/integrals)
            KE_density_local.append(1 / (2 * m) * x.shape[1] * gradient[:, 0] ** 2 / integrals)
            del x_tmp, x_tmp_numerator, log_Psi_numerator_val, x_int
            # del log_Psi_integrand_vals, integrand_vals, integrals, second_ders
            del log_Psi_integrand_vals, integrand_vals, integrals, gradient

        num_density_local = torch.cat(num_density_local)[:, None]
        KE_density_local = torch.cat(KE_density_local)[:, None]
        return num_density_local, KE_density_local

    def IE_density_function(self, N_pts, x_sorted, chain_idx_sorted, n_chains,
                            GPU_device, m, g, N_int, batch_size):
        '''
        Estimates the Lieb-Liniger interaction energy density function
        across the range [0,L], and the standard deviation thereof

        Parameters:
          N_pts: number of points at which the density is estimated over
            the range [0,L]
          x_sorted: a list of tensors, the nth of which contains the sampled
            particle configurations of n particles and is of dimension (#, n),
            where # denotes the number of samples drawn of particle number n
          chain_idx_sorted: a list of tensors denoting the MCMC chain from
            which the corresponding particle configuration in x_sorted came.
          n_chains: number of MCMC chains
          GPU_device: the GPU device under use (i.e. cuda)
          N_int: number of integration points used to compute the numerical integral
            needed for local density calculation
          batch_size: the size of batches over which the local density is
            estimated; performing the numerical integral is memory intensive, and
            hence we batch over entries to save memory

        Returns:
          IE_densities: the estimated Lieb-Liniger interaction energy densities
          IE_densities_std: the standard deviation of the estimated Lieb-Liniger
            interaction energy densities
        '''

        # Evaluates interaction energy density across N_pts points
        xs = np.linspace(0, self.L, N_pts)
        IE_densities = np.zeros(N_pts)
        IE_densities_std = np.zeros(N_pts)
        for i in range(N_pts):
            print('Evaluating interaction energy density at point ' + str(i + 1) + \
                  '/' + str(N_pts))
            x_eval = xs[i]
            IE_density, IE_density_std = \
                self.IE_density(x_eval, x_sorted, chain_idx_sorted, n_chains, \
                                GPU_device, m, g, N_int, batch_size)
            IE_densities[i] = IE_density
            IE_densities_std[i] = IE_density_std

        return IE_densities, IE_densities_std

    def IE_density(self, x_eval, x_sorted, chain_idx_sorted, n_chains,
                   GPU_device, m, g, N_int, batch_size):
        '''
        Estimates the interaction energy density at x_eval and their standard
        deviations

        Parameters:
          x_eval: a scalar value of the point at which the interaction energy
            density is to be estimated
          x_sorted: a list of tensors, the nth of which contains the sampled
            particle configurations of n particles and is of dimension (#, n),
            where # denotes the number of samples drawn of particle number n
          chain_idx_sorted: a list of tensors denoting the MCMC chain from
            which the corresponding particle configuration in x_sorted came.
          n_chains: number of MCMC chains
          GPU_device: the GPU device under use (i.e. cuda)
          N_int: number of integration points used to compute the numerical integral
            needed for local density calculation
          batch_size: the size of batches over which the local density is
            estimated; performing the numerical integral is memory intensive, and
            hence we batch over entries to save memory

        Returns:
          IE_density: the estimated interaction energy density
          IE_density_std: the standard deviation of the estimated interaction
            energy density
        '''

        # Calculates local interaction energy densities
        Local_IE_densities = []
        for x in x_sorted:
            IE_density_local = self.Local_IE_density(x_eval, x, GPU_device, m, g,
                                                     N_int, batch_size)
            Local_IE_densities.append(IE_density_local)

        # Estimated interaction energy density
        Local_IE_densities = torch.cat(Local_IE_densities)
        IE_density = torch.mean(Local_IE_densities)

        # Calculates standard deviations by binning across the MCMC chains
        chain_idx_sorted_flat = torch.cat(chain_idx_sorted)
        IE_density_means_chains = \
            torch.tensor([torch.sum(Local_IE_densities[:, 0] * (chain_idx_sorted_flat == i)) / \
                          torch.sum(chain_idx_sorted_flat == i) \
                          for i in range(n_chains)], device=GPU_device)
        IE_density_std = torch.std(IE_density_means_chains) / (n_chains ** 0.5)

        return IE_density, IE_density_std

    def Local_IE_density(self, x_eval, x, GPU_device, m, g, N_int, batch_size):
        '''
        Calculates the local interaction energy densities of the configurations in x

        Parameters:
          x_eval: a scalar value of the point at which the density is to be estimated
          x: a tensor of dimension (n_samples, n_particles) representing
            n_samples samples of n_particles particle positions
          GPU_device: the GPU device under use (i.e. cuda)
          m: particle mass
          g: Lieb-Liniger interaction strength
          N_int: number of integration points used to compute the numerical integral
            for the marginal distribution
          batch_size: the size of batches over which the local density is estimated;
            performing the numerical integral is memory intensive, and hence we batch
            over entries to save memory

        Returns:
          IE_density_local: a tensor of dimension (n_samples, 1) containing the
          local interaction energy densities at x_eval of the configurations in x
        '''

        x_eval_tensor = torch.cuda.FloatTensor(1).fill_(x_eval)
        # The points across which the numerical integral is calculated
        x_int_points = torch.linspace(0, self.L, N_int, device=GPU_device)

        IE_density_local = []
        for i in range(np.ceil(x.shape[0] / batch_size).astype(int)):
            # Computes log_Psi for numerator
            x_tmp = x[i * batch_size:(i + 1) * batch_size, :]
            x_tmp_numerator = torch.cat((x_eval_tensor.repeat(x_tmp.shape[0], 1),
                                         x_eval_tensor.repeat(x_tmp.shape[0], 1),
                                         x_tmp[:, 2:]), axis=1)
            log_Psi_numerator_val = self.log_Psi(x_tmp_numerator, m, g)[:, 0].detach()

            # Computes log_Psi for integral over distribution
            x_int = torch.cat((x_int_points[:, None, None, None].repeat(1, N_int,
                                                                        x_tmp.shape[0], 1),
                               x_int_points[None, :, None, None].repeat(N_int, 1,
                                                                        x_tmp.shape[0], 1),
                               x_tmp[None, None, :, 2:].repeat(N_int, N_int, 1, 1)), axis=3)
            x_int = x_int.reshape(N_int ** 2 * x_tmp.shape[0], -1)
            log_Psi_integrand_vals = self.log_Psi(x_int, m, g).detach()
            log_Psi_integrand_vals = log_Psi_integrand_vals.reshape(N_int, N_int,
                                                                    x_tmp.shape[0], 1)

            # Calculates local interaction energy density
            integrand_vals = torch.exp(2 * log_Psi_integrand_vals[:, :, :, 0])
            integrals = torch.trapz(torch.trapz(integrand_vals,
                                                dx=self.L / (N_int - 1), dim=0),
                                    dx=self.L / (N_int - 1), dim=0)
            IE_density_local.append(g * x.shape[1] * (x.shape[1] - 1) * \
                                    torch.exp(2 * log_Psi_numerator_val) / integrals)
            del x_tmp, x_tmp_numerator, log_Psi_numerator_val, x_int
            del log_Psi_integrand_vals, integrand_vals, integrals

        IE_density_local = torch.cat(IE_density_local)[:, None]
        return IE_density_local

    def LL_GS_Energy(self, m, g, L, n):
        '''
        Determines the ground state energy of the Lieb-Liniger model in a box (at
        fixed particle number) by solving the transcendental equations outlined in
        "Boundary Energy of a Bose Gas in One Dimension" by M. Gaudin

        Parameters:
          m: particle mass
          g: Lieb-Liniger interaction strength
          L: system size
          n: number of particles

        Returns:
          E: Ground state energy
        '''

        c = 2 * m * g
        n_alpha = np.ones(n)  # The ground state occurs when n_alpha = 1 for all alpha

        # Functions defining the transcendental equations
        def funcs(k_alpha):
            k_alpha = np.array(k_alpha)
            difs = (k_alpha[:, None] - k_alpha[None, :])
            difs += 1e-20 * np.identity(k_alpha.shape[0])
            sums = (k_alpha[:, None] + k_alpha[None, :])
            sum_vals = np.arctan(c / difs) + np.arctan(c / sums)
            np.fill_diagonal(sum_vals, 0.)

            f = k_alpha * L - np.pi * n_alpha - np.sum(sum_vals, axis=1)
            f = list(f)
            return f

        # Solve transcendental equations for k_alpha; energy is 1/(2*m)*sum(k_alpha**2)
        sol = root(funcs, np.pi * np.arange(1, n + 1) / L)
        k_alpha = sol.x
        E = 1 / (2 * m) * np.sum(k_alpha ** 2)

        return E

    def Exact_energy_and_n(self, m, mu, g, L):
        '''
        Determines the exact ground state energy and particel number of the
        Lieb-Liniger model in a box (in the grand canonical ensemble)

        Parameters:
          m: particle mass
          mu: chemical potential
          g: Lieb-Liniger interaction strength
          L: system size
          n: number of particles

        Returns:
          E: Ground state energy
          n: Ground state particle number
        '''

        E = 0
        n = 0

        # Scans through particle number sectors until the ground state is located
        while 1 == 1:
            n_tmp = n + 1
            E_tmp = self.LL_GS_Energy(m, g, L, n_tmp) - mu * n_tmp

            if E_tmp < E:
                E = E_tmp
                n = n_tmp
            else:
                break

        return E, n

    def TG_n(self, L, m, mu):
        '''
        Calculates the particle number of the Tonks-Girardeau ground state

        Parameters:
          L: system size
          m: particle mass
          mu: chemical potential

        Returns:
          n: particle number of the Tonks-Girardeau ground state
        '''

        b = np.pi ** 2 / (2 * m * L ** 2)
        c = np.pi ** 2 / (12 * m * L ** 2) - mu
        n = (-b + (b ** 2 - 4 * b * c) ** 0.5) / (2 * b)
        n = int(np.round(n))
        return n

    def TG_energy(self, L, m, mu):
        '''
        Calculates the energy of the Tonks-Girardeau ground state

        Parameters:
          L: system size
          m: particle mass
          mu: chemical potential

        Returns:
          E: energy of the Tonks-Girardeau ground state
        '''

        n = self.TG_n(L, m, mu)
        b = np.pi ** 2 / (2 * m * L ** 2)
        c = np.pi ** 2 / (12 * m * L ** 2) - mu
        E = b * (n ** 3 / 3 + n ** 2 / 2 + n / 6) - mu * n
        return E

    def TG_number_density(self, x_eval, n, L):
        '''
        Calculates the number density of the Tonks-Girardeau ground state at the
        point x_eval

        Parameters:
          x_eval: a scalar value at which the number density is to be evaluated
          n: particle number
          L: system size

        Returns:
          num_density_TG: an array of the number density
        '''

        num_density_TG = np.sum(np.array([2 / L * (np.sin(j * np.pi * x_eval / L)) ** 2 \
                                          for j in range(1, n + 1)]))
        return num_density_TG

    def TG_KE_density(self, x_eval, n, L, m):
        '''
        Calculates the kinetic energy density of the Tonks-Girardeau ground state
        at the point x_eval

        Parameters:
          x_eval: a scalar value at which the KE density is to be evaluated
          n: particle number
          L: system size
          m: particle mass

        Returns:
          KE_density_TG: an array of the kinetic energy density
        '''

        KE_density_TG = 1 / (2 * m) * np.sum(np.array([2 / L * (j * np.pi / L * np.cos(j * np.pi * x_eval / L)) ** 2 \
                                                       for j in range(1, n + 1)]))
        return KE_density_TG

    def TG_density_function(self, N_pts, n, L, m):
        '''
        Calculates the number density and kinetic energy density of the
        Tonks-Girardeau ground state across the range [0,L]

        Parameters:
          N_pts: number of points at which the number density is estimated over
            the range [0,L]
          n: particle number
          L: system size
          m: particle mass

        Returns:
          num_density_TG: an array of the kinetic energy density
          KE_density_TG: an array of the kinetic energy density
        '''

        xs = np.linspace(0, L, N_pts)
        num_densities_TG = np.array([self.TG_number_density(x, n, L) for x in xs])
        KE_densities_TG = np.array([self.TG_KE_density(x, n, L, m) for x in xs])
        return num_densities_TG, KE_densities_TG

    def plot_number_density(self, num_densities, num_densities_std,
                            num_densities_exact):
        '''
        Plots particle number density vs position

        Parameters:
          num_densities: an array of the approximate particle number densities
          num_densities_std: an array of the standard deviation of the approximate
            particle number densities
          num_densities_exact: an array of the exact particle number densities
        '''

        # Plots particle number density
        xs = np.linspace(0, self.L, num_densities.shape[0])
        fig, ax = plt.subplots(figsize=(3.25, 2.8), dpi=200)
        ax.plot(xs, num_densities, linewidth=1, label='Approximate')
        ax.fill_between(xs, num_densities - num_densities_std,
                        num_densities + num_densities_std, facecolor='blue',
                        alpha=0.25, interpolate=True)
        ax.plot(xs, num_densities_exact, linewidth=1, label='Exact')
        ax.set_title('Particle Number Density vs. x', fontsize=10)
        ax.set_ylabel('Number Density', fontsize=8)
        ax.set_xlabel('x', fontsize=8)
        ax.legend(prop={'size': 7})
        ax.set_xticklabels(np.round(ax.get_xticks(), 2), fontsize=6)
        ax.set_yticklabels(np.round(ax.get_yticks(), 2), fontsize=6)

        fig.tight_layout()
        fig.show()

    def plot_KE_density(self, KE_densities, KE_densities_std,
                        KE_densities_exact):
        '''
        Plots kinetic energy density vs position

        Parameters:
          KE_densities: an array of the approximate kinetic energy densities
          KE_densities_std: an array of the standard deviation of the approximate
            kinetic energy densities
          KE_densities_exact: an array of the exact kinetic energy densities
        '''

        # Plots kinetic energy density
        xs = np.linspace(0, self.L, KE_densities.shape[0])
        fig, ax = plt.subplots(figsize=(3.25, 2.8), dpi=200)
        ax.plot(xs, KE_densities, linewidth=1, label='Approximate')
        ax.fill_between(xs, KE_densities - KE_densities_std,
                        KE_densities + KE_densities_std, facecolor='blue',
                        alpha=0.25, interpolate=True)
        ax.plot(xs, KE_densities_exact, linewidth=1, label='Exact')
        ax.set_title('Kinetic Energy Density vs. x', fontsize=10)
        ax.set_ylabel('KE Density', fontsize=8)
        ax.set_xlabel('x', fontsize=8)
        ax.legend(prop={'size': 7})
        ax.set_xticklabels(np.round(ax.get_xticks(), 2), fontsize=6)
        ax.set_yticklabels(np.round(ax.get_yticks(), 2), fontsize=6)

        fig.tight_layout()
        fig.show()

    def plot_IE_density(self, IE_densities, IE_densities_std):
        '''
        Plots interaction energy density vs position

        Parameters:
          IE_densities: an array of the approximate interaction energy densities
          IE_densities_std: an array of the standard deviation of the approximate
            interaction energy densities
        '''

        # Plots interaction energy density
        xs = np.linspace(0, self.L, IE_densities.shape[0])
        fig, ax = plt.subplots(figsize=(3.25, 2.8), dpi=200)
        ax.plot(xs, IE_densities, linewidth=1, label='Approximate')
        ax.fill_between(xs, IE_densities - IE_densities_std,
                        IE_densities + IE_densities_std, facecolor='blue',
                        alpha=0.25, interpolate=True)
        ax.plot(xs, 0 * xs, linewidth=1, label='Exact')  # Exact IE density is 0
        ax.set_title('Interaction Energy Density vs. x', fontsize=10)
        ax.set_ylabel('Interaction Energy Density', fontsize=8)
        ax.set_xlabel('x', fontsize=8)
        ax.legend(prop={'size': 7})
        ax.set_xticklabels(np.round(ax.get_xticks(), 2), fontsize=6)
        ax.set_yticklabels(np.round(ax.get_yticks(), 4), fontsize=6)

        fig.tight_layout()
        fig.show()

    def plot_E_density(self, E_densities, E_densities_std, E_densities_exact):
        '''
        Plots interaction energy density vs. position

        Parameters:
          E_densities: an array of the approximate total energy densities
          E_densities_std: an array of the standard deviation of the approximate
            total energy densities
          E_densities_exact: an array of the exact kinetic energy densities
        '''

        # Plots energy density
        xs = np.linspace(0, self.L, E_densities.shape[0])
        fig, ax = plt.subplots(figsize=(3.25, 2.8), dpi=200)
        ax.plot(xs, E_densities, linewidth=1, label='Approximate')
        ax.fill_between(xs, E_densities - E_densities_std,
                        E_densities + E_densities_std, facecolor='blue',
                        alpha=0.25, interpolate=True)
        ax.plot(xs, E_densities_exact, linewidth=1, label='Exact')
        ax.set_title('Energy Density vs. x', fontsize=10)
        ax.set_ylabel('Energy Density', fontsize=8)
        ax.set_xlabel('x', fontsize=8)
        ax.legend(prop={'size': 7})
        ax.set_xticklabels(np.round(ax.get_xticks(), 2), fontsize=6)
        ax.set_yticklabels(np.round(ax.get_yticks(), 2), fontsize=6)

        fig.tight_layout()
        fig.show()


# 3주차

# 정렬 - 1
# 버블 정렬
# input = [4, 6, 2, 9, 1]

# def bubble_sort(array):
#     # 이 부분을 채워보세요!
#     n = len(array)
#     for i in range(n - 1):         # 4번만 비교하면 되니까!
#             for j in range(n - i - 1): # Q. 여기서 왜 5 - i - 1 일까요? => i의 다음 자리?
#                 if array[j] > array[j + 1]:
#                     array[j], array[j + 1] = array[j + 1], array[j] #Swap a, b = b, a

#     return array

# bubble_sort(input)
# print(input)  # [1, 2, 4, 6, 9] 가 되어야 합니다!

# print("정답 = [1, 2, 4, 6, 9] / 현재 풀이 값 = ",bubble_sort([4, 6, 2, 9, 1]))
# print("정답 = [-1, 3, 9, 17] / 현재 풀이 값 = ",bubble_sort([3,-1,17,9]))
# print("정답 = [-3, 32, 44, 56, 100] / 현재 풀이 값 = ",bubble_sort([100,56,-3,32,44]))

##########################
# 정렬 - 2
# 선택 정렬 => 최솟값을 찾아서 변경
# input = [4, 6, 2, 9, 1]

# def insertion_sort(array):
#     # 이 부분을 채워보세요!
#     n = len(array)
#     for i in range(n - 1):   #맨마지막에 하나 남은 원소를 비교하지 않음
#         min_index = i
#         for j in range(n - i): #점점 줄어들기 때문에 i가 줄어듬
#             if array [i + j] < array[min_index]:
#                 min_index = i + j

#         array[i], array[min_index] = array[min_index], array[i]

#     return array

# insertion_sort(input)
# print(input) # [1, 2, 4, 6, 9] 가 되어야 합니다!

# 삽입 정렬 => 전체에서 하나씩 올바른 위치에 삽입
# input = [4, 6, 2, 9, 1]

# def insertion_sort(array):
#     # 이 부분을 채워보세요!
#     n = len(array)
#     for i in range(1, n):
#         for j in range(i):
#             if array[i - j - 1] > array[i - j]: #앞에있는 애가 뒤에있는 애보다 크다면 변경!
#                 array[i - j - 1], array[i - j] = array[i - j], array[i - j - 1]
#             else:
#                 break #더 작지 않을 경우에는 반복문 실행 X => ex.9

#     return

# insertion_sort(input)
# print(input) # [1, 2, 4, 6, 9] 가 되어야 합니다!

# print("정답 = [4, 5, 7, 7, 8] / 현재 풀이 값 = ",insertion_sort([5,8,4,7,7]))
# print("정답 = [-1, 3, 9, 17] / 현재 풀이 값 = ",insertion_sort([3,-1,17,9]))
# print("정답 = [-3, 32, 44, 56, 100] / 현재 풀이 값 = ",insertion_sort([100,56,-3,32,44]))

##########################
# 정렬 - 3
# 병합 정렬 - merge
# array_a = [1, 2, 3, 5]
# array_b = [4, 6, 7, 8]

# def merge(array1, array2):
#     # 이 부분을 채워보세요!
#     array_c = []
#     array1_index = 0
#     array2_index = 0
#     while array1_index < len(array1) and array2_index < len(array2): #각 array 의 길이까지
#         if array1[array1_index] < array2[array2_index]: #각 index 값을 비교
#             array_c.append(array1[array1_index])    #c에다가 append
#             array1_index += 1 #다음것으로
#         else:
#             array_c.append(array2[array2_index])
#             array2_index += 1

#     if array1_index == len(array1): #array1이 끝까지 갔다 == array2가 남아있다
#         while array2_index < len(array2):
#             array_c.append(array2[array2_index])
#             array2_index += 1

#     if array2_index == len(array2):
#         while array1_index < len(array1):
#             array_c.append(array1[array1_index])
#             array1_index += 1
#     return array_c

# print(merge(array_a, array_b))  # [1, 2, 3, 4, 5, 6, 7, 8] 가 되어야 합니다!

# 병합 정렬 - mergeSort =>: 할 정복의 개념
# MergeSort(시작점, 끝점)

# 그러면
# MergeSort(0, N) = Merge(MergeSort(0, N/2) + MergeSort(N/2, N))

# array = [5, 3, 2, 1, 6, 8, 7, 4]

# def merge_sort(array):
#     # 이 곳을 채워보세요!
#     if len(array) <= 1:
#         return array
#     mid = len(array) // 2
#     left_array = merge_sort(array[:mid])   # 왼쪽 부분을 정렬하고 => :mid 0부터 mid 까지
#     right_array = merge_sort(array[mid:])  # 오른쪽 부분을 정렬한 다음에
#     merge(left_array, right_array)         # 합치면서 정렬
#     return merge(merge_sort(left_array), merge_sort(right_array))
#     #return array

# def merge(array1, array2):
#     result = []
#     array1_index = 0
#     array2_index = 0
#     while array1_index < len(array1) and array2_index < len(array2):
#         if array1[array1_index] < array2[array2_index]:
#             result.append(array1[array1_index])
#             array1_index += 1
#         else:
#             result.append(array2[array2_index])
#             array2_index += 1

#     if array1_index == len(array1):
#         while array2_index < len(array2):
#             result.append(array2[array2_index])
#             array2_index += 1

#     if array2_index == len(array2):
#         while array1_index < len(array1):
#             result.append(array1[array1_index])
#             array1_index += 1

#     return result

# print(merge_sort(array))  # [1, 2, 3, 4, 5, 6, 7, 8] 가 되어야 합니다!

##########################
# 스택 => 한쪽 끝으로만 자료를 넣고 뺄 수 있는 자료 구조. ex) 빨래 바구니
# 스택의 구현
class Node:
    def __init__(self, data):
        self.data = data
        self.next = None

class Stack:
    def __init__(self):
        self.head = None

    def push(self, value): # 현재 [4] 밖에 없다면 #맨 앞에 데이터 넣기
        # 어떻게 하면 될까요?
        new_head = Node(value)             # [3] 을 만들고!
        new_head.next = self.head          # [3] -> [4] 로 만든다음에
        self.head = new_head               # 현재 head의 값을 [3] 으로 바꿔준다.
        return

    # pop 기능 구현
    def pop(self): #맨 앞의 데이터 뽑기
        # 어떻게 하면 될까요?
        if self.is_empty():                  # 만약 비어있다면 에러!
            return "Stack is empty!"
        delete_head = self.head
        self.head = self.head.next
        return delete_head

    def peek(self): #맨 앞의 데이터 보기
        # 어떻게 하면 될까요?
        if self.is_empty():                  # 만약 비어있다면 에러!
            return "Stack is empty!"
        return self.head.data
        return

    # isEmpty 기능 구현
    def is_empty(self): #스택이 비었는지 안 비었는지 여부 반환해주기
        # 어떻게 하면 될까요?
        return self.head is None

# ex)
stack = Stack()
stack.push(3)
print(stack.peek())
stack.push(4)
print(stack.peek())
print(stack.pop().data)
print(stack.peek())
print(stack.is_empty())

# 문제 풀기 - 탑 (어렵네;)
top_heights = [6, 9, 5, 7, 4]

def get_receiver_top_orders(heights):
    answer = [0] * len(heights)
    while heights:  # heights가 빈상태가 아닐때까지
        height = heights.pop()  # hegihts에서 맨앞
        # heights의 현재 길이 첫 -1 => index번호이기 때문에, 마지막 -1은 0꺼자 -1씩 줄이기
        for idx in range(len(heights) - 1, 0, -1):
            if heights[idx] > height:
                answer[len(heights)] = idx + 1  # +1의 이유는 index가 아닌 위치를 알려주기 위함
                break
    return answer

print(get_receiver_top_orders(top_heights))  # [0, 0, 2, 2, 4] 가 반환되어야 한다!

print("정답 = [0, 0, 2, 2, 4] / 현재 풀이 값 = ",get_receiver_top_orders([6, 9, 5, 7, 4]))
print("정답 = [0, 0, 0, 3, 3, 3, 6] / 현재 풀이 값 = ",get_receiver_top_orders([3, 9, 9, 3, 5, 7, 2]))
print("정답 = [0, 0, 2, 0, 0, 5, 6] / 현재 풀이 값 = ",get_receiver_top_orders([1, 5, 3, 6, 7, 6, 5]))

##########################
# 큐 => 한쪽 끝으로 자료를 넣고, 반대쪽에서는 자료를 뺄 수 있는 선형구조. ex) 놀이기구줄
# ===> 순서대로 처리되어야 하는 일에 필요

class Node:
    def __init__(self, data):
        self.data = data
        self.next = None

class Queue:
    def __init__(self):
        self.head = None
        self.tail = None

    def enqueue(self, value):  # 맨 뒤에 데이터 추가하기
        # 어떻게 하면 될까요?
        new_node = Node(value)
        if self.is_empty():                # 만약 비어있다면,
            self.head = new_node           # head 에 new_node를
            self.tail = new_node           # tail 에 new_node를 넣어준다.
            return

        self.tail.next = new_node
        self.tail = new_node
        return

    def dequeue(self):  # 맨 앞의 데이터 뽑기
        # 어떻게 하면 될까요?
        if self.is_empty():
            return "Queue is empty!"

        delete_head = self.head
        self.head = self.head.next

        return delete_head

    def peek(self): #맨 앞의 데이터 보기
        # 어떻게 하면 될까요?
        if self.is_empty():
            return "Queue is empty!"

        return self.head.data

    def is_empty(self): #큐가 비었는지 안 비었는지 여부 반환해주기
        # 어떻게 하면 될까요?
        return self.head is None

#ex)
queue = Queue()
queue.enqueue(3) #맨 뒤에 3
print(queue.peek()) #맨 앞을 보면 3뿐이니 3
queue.enqueue(4) #맨 뒤에 4
print(queue.peek()) #맨 앞 3
queue.enqueue(5) #맨 뒤에 5
print(queue.peek()) #맨 앞 3
print(queue.dequeue()) # 맨 앞 데이터 뽑기
print(queue.peek()) #맨 앞 3 뽑혀서 4

##########################
# 해쉬 - 1
# 딕셔너리의 구현
class Dict:
    def __init__(self):
        self.items = [None] * 8

    def put(self, key, value):
        # 구현해보세요!
        index = hash(key) % len(self.items) #현재 배열의 최대 길이로 hash => 인덱스의 범위에 맞춘 뒤 그 값을 담는다
        self.items[index] = value
        return

    def get(self, key):
        # 구현해보세요!
        index = hash(key) % len(self.items)
        return self.items[index]

my_dict = Dict()
my_dict.put("test", 3)
print(my_dict.get("test"))  # 3이 반환되어야 합니다!

# =====> 충돌 링크드 리스트로 딕셔너리로 해결 (추가되는 녀석 뒤에 붙이기)
class LinkedTuple:
    def __init__(self):
        self.items = list()

    def add(self, key, value):
        self.items.append((key, value))

    def get(self, key):
        for k, v in self.items:
            if key == k:
                return v

class LinkedDict:
    def __init__(self):
        self.items = [] #[LinkedTuple(), LinkedTuple(), LinkedTuple()...]
        for i in range(8):
            self.items.append(LinkedTuple())

    def put(self, key, value):
        index = hash(key) % len(self.items)
        self.items[index].add(key, value)

    def get(self, key):
        index = hash(key) % len(self.items)
        return self.items[index].get(key)

##########################
# 해쉬 - 2
# 해쉬 테이블(Hash Table) => "키" 와 "데이터"를 저장함으로써 즉각적으로 데이터를 받아오고 업데이트하고 싶을 때 사용하는 자료구조
# ========================> 딕셔너리와 동일

#해쉬 함수(Hash Function) =>  임의의 길이를 갖는 메시지를 입력하여 고정된 길이의 해쉬값을 출력하는 함수

#만약, 해쉬 값 혹은 인덱스가 중복되어 충돌이 일어난다면? 체이닝과 개방 주소법 방법으로 해결

### 출석체크 문제
# all_students = ["나연", "정연", "모모", "사나", "지효", "미나", "다현", "채영", "쯔위"]
# present_students = ["정연", "모모", "채영", "쯔위", "사나", "나연", "미나", "다현"]
# #정열이 되어있지 않으니 반복문 돌려보기

def get_absent_student(all_array, present_array):
    # 구현해보세요!
    dict = {}
    for key in all_array:
        dict[key] = True #아무 값이나 넣기

    for key in present_array: #dict에서 key를 하나씩 없애기
        del dict[key]

    for key in dict.keys(): #key 중에 하나를 반환
        return key

print(get_absent_student(all_students, present_students))

print("정답 = 예지 / 현재 풀이 값 = ",get_absent_student(["류진","예지","채령","리아","유나"],["리아","류진","채령","유나"]))
print("정답 = RM / 현재 풀이 값 = ",get_absent_student(["정국","진","뷔","슈가","지민","RM"],["뷔","정국","지민","진","슈가"]))

##########################
#숙제

#쓱 최대로 할인 적용하기
shop_prices = [30000, 2000, 1500000]
user_coupons = [20, 40]

def get_max_discounted_price(prices, coupons):
    # 이 곳을 채워보세요!
    prices.sort(reverse=True)
    coupons.sort(reverse=True)
    price_index = 0
    coupon_index = 0
    max_discount_price = 0

    while price_index < len(prices) and coupon_index < len(coupons):
        max_discount_price += prices[price_index] * (100 - coupons[coupon_index]) / 100
        price_index += 1
        coupon_index += 1

    while price_index < len(prices):
        max_discount_price += prices[price_index]
        price_index += 1

    return max_discount_price

print("정답 = 926000 / 현재 풀이 값 = ", get_max_discounted_price([30000, 2000, 1500000], [20, 40]))
print("정답 = 485000 / 현재 풀이 값 = ", get_max_discounted_price([50000, 1500000], [10, 70, 30, 20]))
print("정답 = 1550000 / 현재 풀이 값 = ", get_max_discounted_price([50000, 1500000], []))
print("정답 = 1458000 / 현재 풀이 값 = ", get_max_discounted_price([20000, 100000, 1500000], [10, 10, 10]))

#올바른 괄호
def is_correct_parenthesis(string):
    # 구현해보세요!
    stack = []

    for i in range(len(string)):
        if string[i] == '(':
            stack.append(i)
        elif string[i] == ")":
            if len(stack) == 0:
                return False
            stack.pop()

    if len(stack) != 0:
        return False
    else:
        return True

print("정답 = True / 현재 풀이 값 = ", is_correct_parenthesis("(())"))
print("정답 = False / 현재 풀이 값 = ", is_correct_parenthesis(")"))
print("정답 = False / 현재 풀이 값 = ", is_correct_parenthesis("((())))"))
print("정답 = False / 현재 풀이 값 = ", is_correct_parenthesis("())()"))
print("정답 = False / 현재 풀이 값 = ", is_correct_parenthesis("((())"))

# #멜론 베스트 앨범 뽑기 (제일 어려움;)
def get_melon_best_album(genre_array, play_array):
    n = len(genre_array)
    genre_total_play_dict = {}
    genre_index_play_array_dict = {}
    for i in range(n):
        genre = genre_array[i]
        play = play_array[i]
        if genre not in genre_total_play_dict:
            genre_total_play_dict[genre] = play
            genre_index_play_array_dict[genre] = [[i, play]]
        else:
            genre_total_play_dict[genre] += play
            genre_index_play_array_dict[genre].append([i, play])

    sorted_genre_play_array = sorted(genre_total_play_dict.items(), key=lambda item: item[1], reverse=True)
    result = []
    for genre, _value in sorted_genre_play_array:
        index_play_array = genre_index_play_array_dict[genre]
        sorted_by_play_and_index_play_index_array = sorted(index_play_array, key=lambda item: item[1], reverse=True)
        for i in range(len(sorted_by_play_and_index_play_index_array)):
            if i > 1:
                break
            result.append(sorted_by_play_and_index_play_index_array[i][0])
    return result

print("정답 = [4, 1, 3, 0] / 현재 풀이 값 = ",   get_melon_best_album(["classic", "pop", "classic", "classic", "pop"], [500, 600, 150, 800, 2500]))
print("정답 = [0, 6, 5, 2, 4, 1] / 현재 풀이 값 = ", get_melon_best_album(["hiphop", "classic", "pop", "classic", "classic", "pop", "hiphop"], [2000, 500, 600, 150, 800, 2500, 2000]))


"""Shift string characters by k position. letter should wrap around, z shifts one return a. """

def caesarCipherEncryptor(string, key):
    shiftedString = ""
    for char in string:
        unicode = (ord(char) - ord('a') + key) % 26 + ord('a')
        shiftedString = shiftedString + chr(unicode)
    return shiftedString

def caesarCipherEncryptor2(string, key):
    shifted = []
    newKey = key
    alphabet = list('abcdefghijklmnopqrstuvwxyz')
    for char in string:
        newLetterCode = (alphabet.index(char) + newKey) % 26
        shifted.append(alphabet[newLetterCode])
    return ''.join(shifted)

import unittest

class TestProgram(unittest.TestCase):
    def test_case_1(self):
        self.assertEqual("zab", caesarCipherEncryptor("xyz", 2))
        self.assertEqual("zab", caesarCipherEncryptor2("xyz", 2))


import io
import csv
import logging
import requests
import glob, os
import numpy as np
import pandas as pd
from pathlib import Path
from astropy.table import Table
from astropy import units as u
from astropy.time import Time
from astropy.coordinates import SkyCoord
import matplotlib.pyplot as plt

# ["g", "r", "i", "z"], value=[1, 2, 3, 4]
dic_filters = {1: "g", 2: "r", 3: "i", 4: "z"}

def i_minus_g(fid, mag):
    """ Compute i-g based on vectors of filters and magnitudes
    """

    if list(set(fid)) == [1, 3]:
        # last measurement
        last_fid = fid[-1]

        # last measurement with different filter
        # could take the mean
        index_other = np.where(np.array(fid) != last_fid)[0][-1]

        sign = np.diff([fid[index_other], last_fid])[0]
        mag = [mag[index_other], mag[-1]]

        return -1 * sign * np.diff(mag)[0]

    else:
        return np.nan

def extract_delta_color(pdf: pd.DataFrame, smooth_by: float):
    """ Extract last g-r and delta mag for each object

    Modified from Fink https://github.com/astrolabsoftware/fink-science-portal/apps/utils.py
    Adding a grouping for a given mjd interval

    Parameters
    ----------
    pdf: pandas DataFrame
        DataFrame containing magnitude, mjd, filters
        Filter band, as integer. g: 1, r: 2, i: 3, z:4
    smooth_by: float
        MJD delta to smooth computation

    Returns
    ----------

    """

    # dmag, rate for each filter
    dic_dmag = dict.fromkeys([1, 2, 3, 4])
    dic_dmag_mjd = dict.fromkeys([1, 2, 3, 4])
    dic_rate = dict.fromkeys([1, 2, 3, 4])
    for fil in pdf["filter"].unique():
        subpdf = pdf[pdf["filter"] == fil]
        subpdf = subpdf.sort_values("mjd", ascending=False)

        # dmag smoothed
        min_mjd = float(int(subpdf.mjd.min()))
        max_mjd = float(int(subpdf.mjd.max()))
        bins = np.arange(min_mjd, max_mjd + 1, smooth_by)
        df_grouped = subpdf.groupby(np.digitize(subpdf.mjd, bins)).median()
        mag_grouped = df_grouped["magnitude"]
        mjd_grouped = df_grouped["mjd"]
        if len(mag_grouped) > 1:
            # only compute if more than one observation
            dmag_ = mag_grouped.diff(periods=1).values[1:]
            dmag_mjd_ = mjd_grouped.values[1:].astype(int)
            dic_dmag[fil] = dmag_
            dic_dmag_mjd[fil] = dmag_mjd_

        # Rate by day
        bins = np.arange(min_mjd, max_mjd + 1, 1)
        df_grouped_byday = subpdf.groupby(np.digitize(subpdf.mjd, bins)).median()
        mag_grouped_byday = df_grouped_byday["magnitude"]
        rate_ = mag_grouped_byday.diff(periods=1).values
        if len(mag_grouped_byday) > 1:
            dic_rate[fil] = rate_[1:]

    # for color (can be modified to take median mag)
    # group by night
    gpdf = pdf.groupby("mjd_int")[["filter", "mjd", "magnitude"]].agg(list)
    # take only nights with at least measurements on 2 different filters
    mask = gpdf["filter"].apply(lambda x: (len(x) > 1) & (np.sum(x) / len(x) != x[0]))
    gpdf_night = gpdf[mask]

    # compute i-g for those nights
    color_tmp = [
        i_minus_g(i, j)
        for i, j in zip(gpdf_night["filter"].values, gpdf_night["magnitude"].values)
    ]
    mask = np.argwhere(~np.isnan(color_tmp)).flatten()
    if len(mask) > 1:
        color_mjd = gpdf_night.index.values[mask]
        color = [color_tmp[k] for k in mask]
    else:
        color_mjd = []
        color = []

    return (dic_dmag, dic_dmag_mjd, dic_rate, color, color_mjd)

if __name__ == "__main__":
    inpath = "./S82sub8_59.12"
    list_files = glob.glob(f"{inpath}/*.forced.difflc.txt")

    # If daily cadence then = rate
    smooth_by = 1

    for fname in list_files[:1]:
        # read file and convert to pandas
        df_tmp = Table.read(fname, format="ascii").to_pandas()

        # variable reformatting
        t = Time(df_tmp["dateobs"].to_list(), format="isot", scale="utc")
        df_tmp["mjd"] = t.mjd.astype(float)
        df_tmp["mjd_int"] = t.mjd.astype(int)
        df_tmp["magnitude"] = df_tmp["m"].values
        df_tmp["filter"] = df_tmp["filt"].replace(
            to_replace=["g", "r", "i", "z"], value=[1, 2, 3, 4]
        )
        df_tmp = df_tmp[df_tmp["magnitude"] != "-"]
        df_tmp["magnitude"] = df_tmp["magnitude"].astype(np.float).copy()

        dic_dmag, dic_dmag_mjd, dic_rate, color, color_mjd = extract_delta_color(
            df_tmp, smooth_by=smooth_by
        )

        import ipdb

        ipdb.set_trace()

        # # Plot to verify processing
        # from matplotlib import gridspec

        # fig = plt.figure(figsize=(14, 14))
        # gs = gridspec.GridSpec(2, 1, height_ratios=[2, 1])
        # ax0 = plt.subplot(gs[0, 0])
        # ax1 = plt.subplot(gs[1, 0])
        # x_arr = np.arange(df_tmp.mjd.min(), df_tmp.mjd.max() + 1, 1)
        # ax1.plot(x_arr, np.zeros(len(x_arr)), color="grey", linestyle="-")

        # for fil in df_tmp["filter"].unique():
        #     sel = df_tmp[df_tmp["filter"] == fil]
        #     if len(sel) > 1:
        #         ax0.errorbar(
        #             sel.mjd,
        #             sel.magnitude,
        #             yerr=sel.dm.astype(float),
        #             fmt="o",
        #             label=dic_filters[fil],
        #         )
        #         if fil in dic_dmag_mjd.keys():
        #             ax1.scatter(dic_dmag_mjd[fil], dic_dmag[fil])
        #             if len(color) > 0:
        #                 ax1.scatter(color_mjd, color, color="black")
        # idx = Path(fname).stem.replace(".forced.difflc", "")
        # ax0.legend()
        # plt.savefig(f"{idx}.png")


from collections import deque
import sys

input = sys.stdin.readline

f_inf = float("inf")

n, m = map(int, input().split())
to = [[] for _ in range(n)]

for _ in range(m):
    a, b = map(int, input().split())
    a -= 1
    b -= 1
    to[a].append(b)
    to[b].append(a)

que = deque([])
que.append(0)

dist = [f_inf for _ in range(n)]
dist[0] = 0

guide = [-1 for _ in range(n)]

while que:
    v = que.popleft()
    for u in to[v]:
        if dist[u] != f_inf: continue
        dist[u] = dist[v] + 1
        guide[u] = v
        que.append(u)

print("Yes")
for i in range(1, n):
    ans = guide[i] + 1
    print(ans)

Atcoder/20200720_BBC_40/D - .. (Double Dots).py