import logging
logging.error('vorsehn')
assignments = []

rows = 'ABCDEFGHI'
cols = '123456789'

def cross(A, B):
    "Cross product of elements in A and elements in B."
    return [s+t for s in A for t in B]

# Define boxes of a sudoku
boxes = cross(rows, cols)
# Define row units of a sudoku
row_units = [cross(r, cols) for r in rows]
# Define column units of a sudoku
column_units = [cross(rows, c) for c in cols]
# Define square units of a sudoku
square_units = [cross(rs, cs) for rs in ('ABC','DEF','GHI') for cs in ('123','456','789')]
# Define diagonal units of a sudoku
#diagonal_units = [['A1','B2','C3','D4','E5','F6','G7','H8','I9'],['A9','B8','C7','D6','E5','F4','G3','H2','I1']]
diagonal_units = [[r+c for r,c in zip(rows,cols)], [r+c for r,c in zip(rows,cols[::-1])]]
# Create a list of all units of a sudoku
unitlist = row_units + column_units + square_units + diagonal_units
# Create a dictionary of units of each box
units = dict((s, [u for u in unitlist if s in u]) for s in boxes)
# Create a dictionary of peers of each box
peers = dict((s, set(sum(units[s],[]))-set([s])) for s in boxes)

def assign_value(values, box, value):
    """
    Please use this function to update your values dictionary!
    Assigns a value to a given box. If it updates the board record it.
    """

    # Don't waste memory appending actions that don't actually change any values
    if values[box] == value:
        return values

    values[box] = value
    if len(value) == 1:
        assignments.append(values.copy())
    return values

def naked_twins(values):
    """Eliminate values using the naked twins strategy.
    Args:
        values(dict): a dictionary of the form {'box_name': '123456789', ...}

    Returns:
        the values dictionary with the naked twins eliminated from peers.
    """

    # Find all instances of naked twins
    # Eliminate the naked twins as possibilities for their peers
    for unit in unitlist:
        # create empty list for naked_twins
        twins = []
        #Iterate over each box (field) in one unit
        for box in unit:
            #check number of possible numbers in box
            le = len(values[box])
            if le == 2:
                # Check if the value in one box matches the value of another box in the same unit
                if values[box] in [values[i_box] for i_box in unit if i_box != box]:
                    twins.append(box)
        #if twins found in a unit..
        if len(twins) > 1:
            # for every single number possible in the "twin boxes"
            for single_number in str(values[twins[0]]):
                # replace those numbers in boxes which don#t belong to the twins
                for box in [i_box for i_box in unit if i_box not in twins]:
                    values[box] = values[box].replace(single_number,'')
    return values

def grid_values(grid):
    """
    Convert grid into a dict of {square: char} with '123456789' for empties.
    Args:
        grid(string) - A grid in string form.
    Returns:
        A grid in dictionary form
            Keys: The boxes, e.g., 'A1'
            Values: The value in each box, e.g., '8'. If the box has no value, then the value will be '123456789'.
    """
    assert len(grid) == 81, "Input grid must be a string of length 81 (9x9)"
    values = []
    all_dig = '123456789'
    for c in grid:
        if c == '.':
            values.append(all_dig)
        elif c in all_dig:
            values.append(c)
    return dict(zip(boxes, values))

def display(values):
    """
    Display the values as a 2-D grid.
    Args:
        values(dict): The sudoku in dictionary form
    """
    width = 1+max(len(values[s]) for s in boxes)
    line = '+'.join(['-'*(width*3)]*3)
    for r in rows:
        print(''.join(values[r+c].center(width)+('|' if c in '36' else '')
                      for c in cols))
        if r in 'CF': print(line)
    return

def eliminate(values):
    """Eliminate values from peers of each box with a single value.

    Go through all the boxes, and whenever there is a box with a single value,
    eliminate this value from the set of values of all its peers.

    Args:
        values: Sudoku in dictionary form.
    Returns:
        Resulting Sudoku in dictionary form after eliminating values.
    """
    solved_values = [box for box in values.keys() if len(values[box]) == 1]
    for box in solved_values:
        digit = values[box]
        for peer in peers[box]:
            values[peer] = values[peer].replace(digit,'')
    return values

def only_choice(values):
    """Finalize all values that are the only choice for a unit.

    Go through all the units, and whenever there is a unit with a value
    that only fits in one box, assign the value to this box.

    Input: Sudoku in dictionary form.
    Output: Resulting Sudoku in dictionary form after filling in only choices.
    """
    for unit in unitlist:
        for digit in '123456789':
            dplaces = [box for box in unit if digit in values[box]]
            if len(dplaces) == 1:
                values[dplaces[0]] = digit
    return values

def reduce_puzzle(values):
    stalled = False
    while not stalled:
        # Check how many boxes have a determined value
        solved_values_before = len([box for box in values.keys() if len(values[box]) == 1])

        values = eliminate(values)
        values = only_choice(values)
        values = naked_twins(values)
        # Check how many boxes have a determined value, to compare
        solved_values_after = len([box for box in values.keys() if len(values[box]) == 1])
        # If no new values were added, stop the loop.
        stalled = solved_values_before == solved_values_after
        # Sanity check, return False if there is a box with zero available values:
        if len([box for box in values.keys() if len(values[box]) == 0]):
            return False
    return values

def search(values):
    "Using depth-first search and propagation, try all possible values."
    # First, reduce the puzzle using the previous function
    values = reduce_puzzle(values)
    if values is False:
        return False ## Failed earlier
    if all(len(values[s]) == 1 for s in boxes):
        return values ## Solved!
    # Choose one of the unfilled squares with the fewest possibilities
    n,s = min((len(values[s]), s) for s in boxes if len(values[s]) > 1)
    # Now use recurrence to solve each one of the resulting sudokus, and
    for value in values[s]:
        new_sudoku = values.copy()
        new_sudoku[s] = value
        attempt = search(new_sudoku)
        if attempt:
            return attempt

def solve(grid, diagonal=True):
    """
    Find the solution to a Sudoku grid.
    Args:
        grid(string): a string representing a sudoku grid.
            Example: '2.............62....1....7...6..8...3...9...7...6..4...4....8....52.............3'
    Returns:
        The dictionary representation of the final sudoku grid. False if no solution exists.
    """

    grid = grid_values(grid)
    grid = reduce_puzzle(grid)
    grid = search(grid)

    return grid

if __name__ == '__main__':
    diag_sudoku_grid = ('2.............62....1....7...6..8...3...9...7...6..4...4....8....52.............3')
    display(solve(diag_sudoku_grid))

    try:
        from visualize import visualize_assignments
        visualize_assignments(assignments)

    except SystemExit:
        pass
    except:
        print('We could not visualize your board due to a pygame issue. Not a problem! It is not a requirement.')


import subprocess
import sys
from pathlib import Path

from .console import console

def download_data(database: Path, name: str, bs_params: str = None):
    """Downloads the specific project into the database.

    Args:
        database (Path): path to the seqdat database
        name (str): name/job id of the sequencing project
        bs_params (str, optional): parameters to pass to `bs download project` . Defaults to None.
    """

    cmd = (
        f"bs download project --name {name} -o {database}/{name}/data {bs_params or ''}"
    )
    try:
        with console.status(
            "[hl b] downloading data from basespace",
            spinner="bouncingBar",
            spinner_style="yellow",
        ) as status:  # noqa
            process = subprocess.run(  # noqa
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                shell=True,
                check=True,
                encoding="utf-8",
            )

        console.print("[info]Download finished.")

    except subprocess.CalledProcessError as err:
        console.print("[red]DOWNLOAD FAILED.")
        console.print("[red]see the below output from basespace-cli:")
        print(err.stdout)
        if "could not parse config file" in err.stdout:
            console.print("No config? Have you authenticated bs-cli?")
            console.print("Run the below command to authenticate with basespace")
            console.print("\n bs auth login")
        sys.exit(1)


import torch
import torch.nn as nn

class vgg_layer(nn.Module):
    def __init__(self, nin, nout):
        super(vgg_layer, self).__init__()
        self.main = nn.Sequential(
                nn.Conv2d(nin, nout, 3, 1, 1),
                nn.BatchNorm2d(nout),
                nn.LeakyReLU(0.2, inplace=True)
                )

    def forward(self, input):
        return self.main(input)

class content_encoder(nn.Module):
    def __init__(self, content_dim, nc=1, nf=64):
        super(content_encoder, self).__init__()
        # 64 x 64
        self.c1 = nn.Sequential(
                vgg_layer(nc, 64),
                vgg_layer(64, 64),
                )
        # 32 x 32
        self.c2 = nn.Sequential(
                vgg_layer(64, 128),
                vgg_layer(128, 128),
                )
        # 16 x 16
        self.c3 = nn.Sequential(
                vgg_layer(128, 256),
                vgg_layer(256, 256),
                vgg_layer(256, 256),
                )
        # 8 x 8
        self.c4 = nn.Sequential(
                vgg_layer(256, 512),
                vgg_layer(512, 512),
                vgg_layer(512, 512),
                )
        # 4 x 4
        self.c5 = nn.Sequential(
                nn.Conv2d(512, content_dim, 4, 1, 0),
                nn.BatchNorm2d(content_dim),
                nn.Tanh()
                )
        self.mp = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)

    def forward(self, input):
        h1 = self.c1(input) # 64 -> 32
        h2 = self.c2(self.mp(h1)) # 32 -> 16
        h3 = self.c3(self.mp(h2)) # 16 -> 8
        h4 = self.c4(self.mp(h3)) # 8 -> 4
        h5 = self.c5(self.mp(h4)) # 4 -> 1
        return h5, [h1, h2, h3, h4]

class decoder(nn.Module):
    def __init__(self, content_dim, pose_dim, nc=1, nf=64):
        super(decoder, self).__init__()
        # 1 x 1 -> 4 x 4
        self.upc1 = nn.Sequential(
                nn.ConvTranspose2d(content_dim+pose_dim, 512, 4, 1, 0),
                nn.BatchNorm2d(nf * 8),
                nn.LeakyReLU(0.2, inplace=True)
                )
        # 8 x 8
        self.upc2 = nn.Sequential(
                vgg_layer(512*2, 512),
                vgg_layer(512, 512),
                vgg_layer(512, 256)
                )
        # 16 x 16
        self.upc3 = nn.Sequential(
                vgg_layer(256*2, 256),
                vgg_layer(256, 256),
                vgg_layer(256, 128)
                )
        # 32 x 32
        self.upc4 = nn.Sequential(
                vgg_layer(128*2, 128),
                vgg_layer(128, 64)
                )
        # 64 x 64
        self.upc5 = nn.Sequential(
                vgg_layer(64*2, 64),
                nn.ConvTranspose2d(nf, nc, 3, 1, 1),
                nn.Sigmoid()
                )
        self.up = nn.UpsamplingNearest2d(scale_factor=2)

    def forward(self, input):
        content, pose = input
        content, skip = content
        if type(pose) == list and len(pose) == 0:
            input_vec = content
        else:
            input_vec = torch.cat([content, pose], 1)
        d1 = self.upc1(input_vec) # 1 -> 4
        up1 = self.up(d1) # 4 -> 8
        d2 = self.upc2(torch.cat([up1, skip[3]], 1)) # 8 x 8
        up2 = self.up(d2) # 8 -> 16
        d3 = self.upc3(torch.cat([up2, skip[2]], 1)) # 16 x 16
        up3 = self.up(d3) # 8 -> 32
        d4 = self.upc4(torch.cat([up3, skip[1]], 1)) # 32 x 32
        up4 = self.up(d4) # 32 -> 64
        output = self.upc5(torch.cat([up4, skip[0]], 1)) # 64 x 64
        return output

class scene_discriminator(nn.Module):
  def __init__(self, pose_dim, nf=256):
    super(scene_discriminator, self).__init__()
    self.pose_dim = pose_dim
    self.main = nn.Sequential(
      nn.Linear(pose_dim*2, nf),
      nn.ReLU(True),
      nn.Linear(nf, nf),
      nn.ReLU(True),
      nn.Linear(nf, 1),
      nn.Sigmoid(),
    )

  def forward(self, input):
    output = self.main(torch.cat(input, 1).view(-1, self.pose_dim*2))
    return output


import math
import os

import numpy
import torch
import numpy as np
from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from sklearn import decomposition
from sklearn.preprocessing import StandardScaler
plt.rcParams["font.family"] = "Times New Roman"
from shared_parameters import channel_parameters

if "siso" in os.getcwd():
    from wireless_autoencoder.siso.parameters import model_parameters
else:
    from wireless_autoencoder.mimo.parameters import model_parameters

if model_parameters['device'] == 'auto':
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
else:
    device = torch.device(model_parameters['device'])

def frange(x, y, jump):
    while x < y:
        yield x
        x += jump

def jakes_flat(fd=926, Ts=1e-6, Ns=1, t0=0, E0=np.sqrt(1), phi_N=0):
    """
        :param:
            fd      : Doppler frequency
            Ts      : sampling period
            Ns      : number of samples
            t0      : initial time
            E0      : channel power
            phi_N   : initial phase of the maximum doppler frequency sinusoid

        :return:
            h       : complex fading vector
    """
    N0 = 8
    N = 4 * N0 + 2
    wd = 2 * np.pi * fd
    t = t0 + np.asarray([i for i in range(0, Ns)]) * Ts
    H = np.ones((Ns, 2))
    coff = E0 / np.sqrt(2 * N0 + 1)
    phi_n = np.asarray([np.pi * i / (N0 + 1) for i in range(1, N0 + 1)])
    w_n = np.asarray([wd * np.cos(2 * np.pi * i / N) for i in range(1, N0 + 1)])
    h_i = np.ones((N0 + 1, Ns))
    for i in range(N0):
        h_i[i, :] = 2 * np.cos(phi_n[i]) * np.cos(w_n[i] * t)
    h_i[N0, :] = np.sqrt(2) * np.cos(phi_N) * np.cos(wd * t)
    h_q = np.ones((N0 + 1, Ns))
    for i in range(N0):
        h_q[i, :] = 2 * np.sin(phi_n[i]) * np.cos(w_n[i] * t)
    h_q[N0, :] = np.sqrt(2) * np.sin(phi_N) * np.cos(wd * t)
    h_I = coff * np.sum(h_i, 0)
    h_Q = coff * np.sum(h_q, 0)
    H[:, 0] = h_I
    H[:, 1] = h_Q
    H = torch.from_numpy(H)
    return torch.unsqueeze(torch.view_as_complex(H.type(torch.float32)), dim=1)

def plot_constellation(data, n_channel=model_parameters['n_channel']):
    plts = []
    for i in range(0, n_channel):
        i = 8
        image = plt.figure(figsize=(6, 6))
        for signals, marker, color, size, *index in data:
            if isinstance(signals, torch.Tensor):
                signals = signals.cpu().detach().numpy()
            plt.scatter(signals[:, 0:i], signals[:, i:n_channel + i], marker=marker, c=color, s=size,
                        label=('Autoencoder' + str(index[0]) if len(index) and i == 0 else ""))

        tick_lmt = [-3, 3]
        ticks_lbl_rng = [-3, 4]
        ticks_step = 1
        plt.xlim(tick_lmt)
        plt.ylim(tick_lmt)

        legends = []
        for i in range(len(data)):
            legends.append(f"Autoencoder{i}")
        plt.legend(legends, fontsize=16, loc="lower left")

        image.axes[0].axhline(0, linestyle='-', color='gray', alpha=0.5)
        image.axes[0].axvline(0, linestyle='-', color='gray', alpha=0.5)
        plt.xticks(list(range(ticks_lbl_rng[0], ticks_lbl_rng[1], ticks_step)), fontsize=16)
        plt.yticks(list(range(ticks_lbl_rng[0], ticks_lbl_rng[1], ticks_step)), fontsize=16)
        plt.xlabel('In-Phase', fontsize=20)
        plt.ylabel('Quadrature', fontsize=20)

        plts.append(image)
        break
    return plts

def constellation_diagram(model):
    model = model.to(device)
    model.load()
    model.eval()

    n_samples = 100
    signals = []
    for j in range(model_parameters['m']):
        x = torch.sparse.torch.eye(model_parameters['m']).index_select(dim=0, index=torch.ones(n_samples).int() * j).to(device)
        encode = model.encoder(x)
        encode = model.channel(encode, ebno=10)
        signals.append(encode)

    data = []
    marker = 'o'
    colors = ['tab:red', 'tab:blue', 'tab:green', 'tab:orange']
    size = 20
    for i, signal in enumerate(signals):
        data.append((signal, marker, colors[i], size))
    plot_constellation(data)

def constellations_diagram(models, pca=False):
    with torch.no_grad():
        n_samples = 100
        signals = []
        for i in range(model_parameters['n_user']):
            encodes = []
            for j in range(model_parameters['m']):
                x = torch.sparse.torch.eye(model_parameters['m']).index_select(dim=0, index=torch.ones(n_samples).int() * j).to(device)
                encode = models[i].encoder(x)
                encode = models[i].channel(encode, ebno=10)
                encodes.append(encode.cpu())
            encodes = torch.stack(encodes, dim=0)
            encodes = encodes.view(-1, model_parameters['n_channel'] * 2)
            signals.append(encodes)

        # h = torch.randn((n_samples, model_parameters['n_user'], model_parameters['n_user']),
        #                 dtype=torch.cfloat).to(device)
        # signals = models[0].channel(encodes, ebno=channel_parameters['ebno'], h=h)

        ''' PCA (Dimensionality reduction) - BEGIN '''
        if pca:
            signals = numpy.stack(signals)
            signals = signals.reshape(-1, model_parameters['n_channel'] * 2)
            # scaler = StandardScaler(with_mean=0, with_std=0.5)
            # signals = scaler.fit_transform(signals)
            pca = decomposition.PCA(n_components=2)
            signals = pca.fit_transform(signals)
            signals = signals.reshape(model_parameters['n_user'], -1, 2)
            # signals = numpy.expand_dims(signals, axis=1)
            # tmp_signals = signals_pca
            # signals = []
            # for signal in tmp_signals:
            #     signal = np.expand_dims(signal, axis=0)
            #     signal = np.repeat(signal, 200, axis=0)
            #     signal = torch.from_numpy(signal).to(device)
            #     encode = models[0].channel(signal)
            #     # encode = signal
            #     signals.append(encode)
            # signals = torch.stack(signals)
            # # signals = torch.from_numpy(signals)
            # signals.transpose_(0, 1)
        ''' PCA (Dimensionality reduction) - END '''

        data = []
        markers = ['o', 'x', '<', '>']
        colors = ['tab:red', 'tab:blue', 'tab:green', 'tab:orange']
        size = 20
        for i, signal in enumerate(signals):
            data.append((signal, markers[i], colors[i], size))

        plots = plot_constellation(data, 1) if pca else plot_constellation(data)
        for i, plot in enumerate(plots):
            label = "pca" if pca else "time_" + str(i + 1)
            label = "projected"
            plot.savefig(f"multi_autoencoder_orthogonal_{label}.png", format="png")
            plot.show()

def bler_chart(model, test_ds, manual_seed=None, return_bler=False):
    if channel_parameters['channel_type'] == 'rayleigh':
        ebno_range = list(frange(0, 22, 5))
    elif channel_parameters['channel_type'] == 'rician':
        ebno_range = list(frange(0, 22, 5))
    else:
        ebno_range = list(frange(-4, 9, 0.25))
    ber = [None] * len(ebno_range)

    x, y = test_ds[:][0], test_ds[:][1]
    x = x.to(device)
    for j in range(0, len(ebno_range)):
        with torch.no_grad():
            encoded_signal = model.encoder(x)
            channel_signal = model.channel(encoded_signal, ebno=ebno_range[j])
            decoded_signal = model.demodulate(channel_signal)

            # decoded_signal = model(x)

            ber[j] = torch.sum(decoded_signal.detach().cpu() != y).item() / len(decoded_signal)
            # print('SNR: {}, BER: {:.8f}'.format(ebno_range[j], ber[j]))

    print(ber)
    with open('snr_heatmap.txt', 'a') as f:
        f.write("%s:%s\n" % (channel_parameters['ebno'], ber))

    if return_bler:
        return ber
    # ber = [math.floor(x * 10000) / 10000 if x < 0.000999 else x for x in ber]
    plt.xlim(ebno_range[0], ebno_range[-1])
    # plt.ylim(ber[-1], 0.5)

    plt.plot(ebno_range, ber, '-ko', clip_on=False,
             label="Autoencoder (" + str(model_parameters['n_channel']) + "," + str(model_parameters['k']) + ")" + (
                 (" - seed: " + str(manual_seed)) if (manual_seed is not None) else ""))

    '''
        Tim O’Shea (Autoencoder + RTN) block error rate for Rayleigh fading channel with 3 taps.
    '''
    # plt.plot(ebno_range, [7.6 * 1e-1, 4e-1, 8 * 1e-2, 6 * 1e-3, 4 * 1e-4], '-s')

    plt.yscale('log')
    plt.xlabel("$E_{b}/N_{0}$ (dB)", fontsize=16)
    plt.ylabel('Block Error Rate', fontsize=16)
    plt.xticks(fontsize=14)
    plt.yticks(fontsize=14)
    plt.grid()
    plt.legend(loc="upper right", ncol=1, fontsize=16)
    plt.tight_layout()
    # plt.savefig("results/png/autoencoder_bler_" + channel_parameters['channel_type'] + ".png")
    # plt.savefig("results/eps/autoencoder_bler_" + channel_parameters['channel_type'] + ".eps", format="eps")
    plt.show()

def blers_chart(models, test_dses, decoder_model, manual_seed=None):
    if decoder_model is not None:
        decoder_model.eval()

    if channel_parameters['channel_type'] == 'rayleigh':
        ebno_range = list(frange(0, 22, 5))
    elif channel_parameters['channel_type'] == 'rician':
        ebno_range = list(frange(0, 22, 5))
    else:
        ebno_range = list(frange(-4, 9, 1.5))
    bers = [[None] * len(ebno_range) for _ in range(model_parameters['n_user'])]

    for j in range(0, len(ebno_range)):
        with torch.no_grad():
            if channel_parameters['channel_type'] == 'awgn':
                encodes = []
                signals = []
                for i in range(model_parameters['n_user']):
                    x = test_dses[i][:][0].to(device)
                    encode = models[i].encoder(x)
                    encodes.append(encode)
                    signals.append(models[i].channel(encode, ebno=ebno_range[j]))

                for i in range(model_parameters['n_user']):
                    for t, encode in enumerate(encodes):
                        if t != i:
                            signals[i] += encode

                signals = torch.stack(signals)
                signals = signals.transpose(0, 1)
                signals = signals.contiguous().view(signals.size()[0], -1)

                for i in range(model_parameters['n_user']):
                    # demodulate = models[i].demodulate(signals[i])
                    with torch.no_grad():
                        decode = decoder_model(signals, i)
                        _, demodulate = torch.max(decode, dim=1)

                    bers[i][j] = torch.sum(demodulate.detach().cpu() != test_dses[i][:][1]).item() / len(demodulate)
            else:
                encodes = []
                for i in range(model_parameters['n_user']):
                    x = test_dses[i][:][0].to(device)
                    encode = models[i].encoder(x)
                    encodes.append(encode)

                encodes = torch.stack(encodes)
                encodes.transpose_(0, 1)

                if channel_parameters['channel_type'] == 'rayleigh':
                    h = torch.randn((x.size()[0], model_parameters['n_user'], model_parameters['n_user']), dtype=torch.cfloat).to(device)
                else:
                    std = (1 / (channel_parameters['channel_k'] + 1)) ** 0.5
                    mean = (channel_parameters['channel_k'] / 2 * (channel_parameters['channel_k'] + 1)) ** 0.5
                    h = torch.normal(mean, std, (x.size()[0], model_parameters['n_user'], model_parameters['n_user']), dtype=torch.cfloat).to(device)

                # signals = []
                # for i in range(model_parameters['n_user']):
                #     signals.append(models[i].channel(encodes, ebno=ebno_range[j], h=h)[:, i, :])
                # signals = torch.stack(signals).transpose(1, 0)
                signals = models[0].channel(encodes, ebno=ebno_range[j], h=h)

                signals = signals.view((-1, model_parameters['n_user'], model_parameters['n_channel'], 2))
                signals = torch.view_as_complex(signals)
                transform = torch.inverse(h) @ signals
                signals = torch.view_as_real(transform).view(transform.size()[0], -1)

                for i in range(model_parameters['n_user']):
                    # demodulate = models[i].demodulate(signals, torch.view_as_real(h).view(-1, 2 * model_parameters[
                    #     'n_user'] * model_parameters['n_user']))
                    with torch.no_grad():
                        # decode = decoder_model(signals, i, torch.view_as_real(h).view(-1, 2 * model_parameters['n_user'] * model_parameters['n_user']))
                        decode = decoder_model(signals, i)
                        _, demodulate = torch.max(decode, dim=1)
                    bers[i][j] = torch.sum(demodulate.detach().cpu() != test_dses[i][:][1]).item() / len(demodulate)

    avg_ber = np.array(bers)
    avg_ber = avg_ber.mean(axis=0)
    avg_ber = avg_ber.tolist()
    print(avg_ber)
    # exit(1)

    if True:
        '''
            Autoencoder saved results for demonstration purposes
        '''
        if channel_parameters['channel_type'] == 'rayleigh':
            bers = [[], [], [], [], [], []]
            bers[0] = [0.33905273437500005, 0.148203125, 0.0540625, 0.017705078125000002, 0.005615234375]
            bers[1] = [0.363525390625, 0.1622216796875, 0.05802734375, 0.0195068359375, 0.006201171875]
            bers[2] = [0.1451, 0.06388, 0.02278, 0.00746, 0.00248]
            bers[3] = [0.26044, 0.11938, 0.0413, 0.0135, 0.00454]
            bers[4] = [0.09048828125, 0.0356640625, 0.01224609375, 0.00390625, 0.0012890625]
            bers[5] = [0.169453125, 0.06681640625, 0.02337890625, 0.0073046875, 0.00232421875]
        elif channel_parameters['channel_type'] == 'rician':
            # bers = [[], [], [], [], [], []]
            bers = [[], [], [], []]
            bers[0] = [0.067451171875, 0.010107421875, 0.002109375, 0.000517578125, 0.000166015625]
            bers[1] = [0.1316650390625, 0.0232080078125, 0.0035205078125, 0.0008935546875, 0.0003125]
            bers[2] = [0.04341796875, 0.01439453125, 0.0040234375, 0.00140625, 0.00041015625]
            bers[3] = [0.08025390625, 0.02673828125, 0.00830078125, 0.00265625, 0.0008984375]
            # bers[4] = [0.082109375, 0.01890625, 0.00310546875, 0.00080078125, 0.00017578125]
            # bers[5] = [0.15662109375, 0.0340234375, 0.00615234375, 0.00166015625, 0.0003515625]
        else:
            bers = [[], [], [], []]
            bers[0] = [0.294208984375, 0.202646484375, 0.118037109375, 0.062529296875, 0.0254296875, 0.008779296875, 0.002001953125, 0.000439453125, 1.953125e-05]
            bers[1] = [0.28151855468750003, 0.189130859375, 0.11747070312499999, 0.061528320312499996, 0.0279150390625, 0.010634765625, 0.0032324218750000005, 0.0006982421875, 0.00015136718750000003]
            bers[2] = [0.1866, 0.14342, 0.10414, 0.06588, 0.03672, 0.01672, 0.00622, 0.00114, 0.00022]
            bers[3] = [0.34154, 0.27248, 0.202, 0.13278, 0.07924, 0.03576, 0.01344, 0.00296, 0.0003]

        bers = [[], [], []]
        bers[0] = [0.708232421875, 0.518984375, 0.328671875, 0.20275390625, 0.139296875]
        bers[1] = [0.69197265625, 0.48836914062500003, 0.292705078125, 0.16257812500000002, 0.09875]
        bers[2] = [0.3118359375, 0.143359375, 0.054541015624999994, 0.018603515625, 0.00611328125]

        line_styles = ['-bo', '-rD', '--k', '-.k', '--g', ':g']
        line_styles = ['-ko', '--rH', '--bH']
        labels = ['2-user', '4-user', 'Simulated BPSK', 'Simulated QPSK', 'Autoencoder BPSK', 'Autoencoder QPSK']
        labels = ['Without Equalization', 'Blind Equalization', 'Zero-Forcing Equalization']
        dash_style = (6, 8)
        # for i in range(2):
        #     labels[i] = "Autoencoder(" + str(model_parameters['n_channel']) + "," + str(
        #         model_parameters['k']) + ")  " + labels[i]
        plt.xlim(ebno_range[0], ebno_range[-1])
        plt.ylim(min(ber[-1] for ber in bers), 0.9)
        for index, ber in enumerate(bers):
            if index == 2 or index == 4:
                extra_args = {'dashes': dash_style}
            else:
                extra_args = {}
            plt.plot(ebno_range, ber, line_styles[index], clip_on=False,
                     label=labels[index], **extra_args)
    else:
        plt.xlim(ebno_range[0], ebno_range[-1])
        # plt.ylim(min(ber[-1] for ber in bers) / 5, 0.5)
        for index, ber in enumerate(bers):
            plt.plot(ebno_range, ber, '-o', clip_on=False,
                     label="Autoencoder" + str(index + 1) + " (" + str(model_parameters['n_channel']) + "," + str(
                         model_parameters['k']) + ")" + (
                               (" - seed: " + str(manual_seed)) if (manual_seed is not None) else ""))

    plt.yscale('log')
    plt.xlabel("$E_{b}/N_{0}$ (dB)", fontsize=18)
    plt.ylabel('Block Error Rate', fontsize=18)
    plt.xticks(fontsize=18)
    plt.yticks(fontsize=18)
    plt.grid()
    plt.legend(loc="lower left", ncol=1, fontsize=15)
    plt.tight_layout()
    plt.savefig("results/png/wireless_autoencoders(" + str(model_parameters['n_channel']) + "," + str(model_parameters['k']) + ")_" + channel_parameters['channel_type'] + ".png")
    plt.savefig("results/eps/wireless_autoencoders(" + str(model_parameters['n_channel']) + "," + str(model_parameters['k']) + ")_" + channel_parameters['channel_type'] + ".eps", format="eps")
    plt.show()
    plt.close()

def losses_chart(losses, multiple=False):
    if not multiple:
        losses = [losses]
    for idx, loss in enumerate(losses, start=1):
        plt.plot(list(range(1, len(loss) + 1)), loss,
                 label="Autoencoder" + (str(idx) if multiple else "") + "(" + str(
                     model_parameters['n_channel']) + "," + str(model_parameters['k']) + ")")
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend(loc="upper right", ncol=1)
    plt.savefig("results/png/losses(" + str(model_parameters['n_channel']) + "," + str(model_parameters['k']) + ").png")
    plt.show()
    plt.close()

def accuracies_chart(accs):
    # plt.figure(figsize=(6, 6))
    fig, ax = plt.subplots()
    plots = ['autoencoder', 'bob', 'willie']
    colors_map = {'autoencoder': 'k', 'bob': 'r', 'willie': 'g'}
    # plots = ['bob', 'willie']
    for plot in plots:
        plt.plot(range(0, len(accs[plot])), accs[plot], colors_map[plot], label=str(plot).capitalize() + " Accuracy")
    # plt.title("Models Accuracies")
    plt.xlabel("Epoch", fontsize=18)
    plt.ylabel("Accuracy", fontsize=18)
    plt.legend(loc="center right", fontsize=16)
    ax.xaxis.set_major_formatter(ticker.FuncFormatter(lambda x, pos: round(x * 10)))
    plt.xticks(fontsize=16)
    plt.yticks(fontsize=16)
    plt.tight_layout()
    plt.savefig("results/png/training_progress_" + channel_parameters['channel_type'] + ".png")
    plt.savefig("results/eps/training_progress_" + channel_parameters['channel_type'] + ".eps", format="eps")
    plt.show()
    plt.close()
    # plt.figure()
    # plt.plot(range(0, len(results['willie_losses'])), results['willie_losses'], label="Willie Losses")
    # plt.xlabel("Epoch")
    # plt.ylabel("Losses")
    # plt.legend(loc="lower right")
    # plt.show()

def accuracy(pred, label):
    return torch.sum(pred == label).item() / len(pred)

def discriminator_accuracy(fake_preds, real_preds):
    fake_preds = fake_preds.cpu()
    real_preds = real_preds.cpu()
    fake_preds[fake_preds > 0.5] = 1
    fake_preds[fake_preds <= 0.5] = 0
    real_preds[real_preds > 0.5] = 1
    real_preds[real_preds <= 0.5] = 0
    discriminator_fake_accuracy = accuracy(fake_preds, torch.zeros(len(fake_preds), 1))
    discriminator_real_accuracy = accuracy(real_preds, torch.ones(len(real_preds), 1))
    return (discriminator_fake_accuracy + discriminator_real_accuracy) / 2

def classifier_accuracy(pred, label):
    _, preds = torch.max(pred, dim=1)
    return accuracy(preds, label)

def reset_grads(*args):
    for arg in args:
        arg.zero_grad()


# -*- coding: utf-8 -*-

from jool.git import BugFilter, FramePopulator, BugMap
from jool.data import Frame
from pygit2 import Repository
import os

def test_clone_repo(gitrepo):
    directory, g = gitrepo
    assert os.path.exists(g.repo_to)

def test_walk_repo(gitrepo):
    directory, g = gitrepo
    g.traverse()
    assert g.dataset.number_of_rows > 0
    assert g.dataset.number_of_columns == 8

def test_filter_has_bugs():
    f = BugFilter()
    bug = f.filter(['this', 'is', 'a', 'bug', 'with', 'fixes'])
    assert bug

def test_filter_no_bugs():
    f = BugFilter()
    bug = f.filter(['this', 'has', 'no', 'issues'])
    assert not bug

def test_populator_extract_key():
    populator = FramePopulator(Frame())

    value = 'commit_id'
    expected = 'id'
    assert populator.extract_key(value) == expected
    value = 'is_bug'
    expected = 'bug'
    assert populator.extract_key(value) == expected

def test_populator_create_transform_classname():
    populator = FramePopulator(Frame())

    value = 'author'
    expected = 'AuthorTransform'
    assert populator.create_transform_classname(value) == expected

def test_blame_mapper(gitrepo):
    directory, g = gitrepo
    repo = Repository("%s/jool/.git" % directory)
    commit = repo.revparse_single('9aeaac94cda9eb7e32d6f861079ab793a0311983')
    mapper = BugMap()
    assert mapper.map(
        repo,
        'c/d.py',
        5,
        commit) == '903b2d2c65f291492f30467edc6442b74a78ab92'


import re
from random import randrange
import allure
from model.contact import Contact

def test_phones_on_home_page(app, db):
    if len(db.get_contact_list()) == 0:
        with allure.step("If contact's list is empty, I'll create a new contact"):
            app.contact.create(Contact(firstname="NAME"))
    with allure.step("Given an  index for compering"):
        index = randrange(len(app.contact.get_contact_list()))
    with allure.step("Given contact's by index information from home page"):
        contact_from_home_page = app.contact.get_contact_list()[index]
    with allure.step("Given contact's by index information from edit page"):
        contact_from_edit_page = app.contact.get_contact_info_from_edit_page(index)
    with allure.step("Then the contact's phones from home page is equal to the contact's phones from edit page"):
        assert clear_number(contact_from_home_page.all_phones_from_home_page) == merge_phones_like_on_home_page(contact_from_edit_page)

def test_phones_on_contact_view_page(app, db):
    if len(db.get_contact_list()) == 0:
        with allure.step("If contact's list is empty, I'll create a new contact"):
            app.contact.create(Contact(firstname="NAME"))
    with allure.step("Given an  index for compering"):
        index = randrange(len(app.contact.get_contact_list()))
    with allure.step("Given contact's by index information from edit page"):
        contact_from_edit_page = app.contact.get_contact_info_from_edit_page(index)
    with allure.step("Given contact's by index information from view page"):
        contact_from_view_page = app.contact.get_contact_info_from_view_page(index)
    with allure.step("Then the contact's phones from edit page is equal to the contact's phones from view page"):
        assert merge_phones_like_on_home_page(contact_from_edit_page) == merge_phones_like_on_home_page(contact_from_view_page)

def clear_number(text):
    return re.sub("[() *+-]", "", text)

def merge_phones_like_on_home_page(contact):
    return "\n".join(filter(lambda x: x != "",  # применяем фильтрацию к списку номеров (убираем пустые номера)
                            map(lambda x: clear_number(x),  # применяем функцию clear_number ко всем номерам
                                filter(lambda x: x is not None,  # применяем фиотрацию к списку номеров (убираем номера равные None)
                                       [contact.home_number, contact.mobile_number, contact.work_number, contact.phone2]))))


import time

import pytest
from fastapi import status

from .conftest import client
from .helpers import create_unique_menu

PREFIX = '/api/v1/menus'

@pytest.mark.asyncio
async def test_get_menus(client):
    get_response = await client.get(f'{PREFIX}/')
    assert get_response.status_code == status.HTTP_200_OK
    assert get_response.json() == []
    menu = create_unique_menu()
    post_response = await client.post(f'{PREFIX}/', json=menu)
    post_data = post_response.json()
    get_response = await client.get(f'{PREFIX}/')
    get_data = get_response.json()
    assert get_response.status_code == status.HTTP_200_OK
    assert get_data[0]['id'] == post_data['id']
    assert get_data[0]['title'] == post_data['title']
    assert get_data[0]['description'] == post_data['description']
    assert isinstance(get_data[0]['submenus_count'], int)
    assert isinstance(get_data[0]['dishes_count'], int)

@pytest.mark.asyncio
async def test_create_menu(client):
    menu = create_unique_menu()
    post_response = await client.post(f'{PREFIX}/', json=menu)
    post_data = post_response.json()
    assert post_response.status_code == status.HTTP_201_CREATED
    assert 'id' in post_data
    assert 'title' in post_data
    assert 'description' in post_data
    assert post_data['title'] == menu['title']
    assert post_data['description'] == menu['description']
    post_response = await client.post(f'{PREFIX}/', json=menu)
    assert post_response.status_code == status.HTTP_409_CONFLICT
    assert post_response.json()['detail'] == 'menu already exists'

@pytest.mark.asyncio
async def test_get_menu_by_id(client):
    menu = create_unique_menu()
    post_response = await client.post(f'{PREFIX}/', json=menu)
    menu_id = post_response.json()['id']
    get_response = await client.get(f'{PREFIX}/{menu_id}')
    get_data = get_response.json()
    assert get_response.status_code == status.HTTP_200_OK
    assert get_data['id'] == menu_id
    assert get_data['title'] == menu['title']
    assert get_data['description'] == menu['description']
    assert isinstance(get_data['submenus_count'], int)
    assert isinstance(get_data['dishes_count'], int)
    menu_id = '0000b000-0e00-00d0-a000-00000be000da'
    get_response = await client.get(f'{PREFIX}/{menu_id}')
    assert get_response.status_code == status.HTTP_404_NOT_FOUND
    assert get_response.json()['detail'] == 'menu not found'

@pytest.mark.asyncio
async def test_update_menu(client):
    menu = create_unique_menu()
    upd_menu = {
        'title': f'updated test menu {str(time.time())}',
        'description': f'updated test menu description {str(time.time())}'
    }
    post_response = await client.post(f'{PREFIX}/', json=menu)
    menu_id = post_response.json()['id']
    patch_response = await client.patch(
        f'{PREFIX}/{menu_id}',
        json=upd_menu
    )
    patch_data = patch_response.json()
    assert patch_response.status_code == status.HTTP_200_OK
    assert patch_data['id'] == menu_id
    assert patch_data['title'] == upd_menu['title']
    assert patch_data['description'] == upd_menu['description']
    menu_id = '0000b000-0e00-00d0-a000-00000be000da'
    patch_response = await client.patch(
        f'{PREFIX}/{menu_id}',
        json=upd_menu
    )
    assert patch_response.status_code == status.HTTP_404_NOT_FOUND
    assert patch_response.json()['detail'] == 'menu not found'

@pytest.mark.asyncio
async def test_delete_menu(client):
    menu = create_unique_menu()
    post_response = await client.post(f'{PREFIX}/', json=menu)
    menu_id = post_response.json()['id']
    del_response = await client.delete(f'{PREFIX}/{menu_id}')
    assert del_response.json()['message'] == 'menu deleted successfully'
    menu_id = '0000b000-0e00-00d0-a000-00000be000da'
    del_response = await client.delete(f'{PREFIX}/{menu_id}')
    assert del_response.status_code == status.HTTP_404_NOT_FOUND
    assert del_response.json()['detail'] == 'menu not found'


import binascii
import ctypes
import struct

class RTP(ctypes.Structure):

  _fields_ = [('flags', ctypes.c_ubyte),      # Flags
              ('mpt',   ctypes.c_ubyte),      # Marker + Payload Type
              ('seq',   ctypes.c_ushort),     # Sequence Number
              ('ts',    ctypes.c_ulong),      # Timestamp
              ('ssrc',  ctypes.c_ulong)]      # Sync Source

  rtp_header_size = 12

  def __init__(self, packet):
    fields = struct.unpack("!BBHII", packet[:self.rtp_header_size])
    self.flags  = fields[0]
    self.marker = fields[1] & 0x1
    self.pt     = (fields[1] & 0xFE) << 1
    self.seq    = fields[2]
    self.ts     = fields[3]
    self.ssrc   = fields[4]
    self.payload = ctypes.c_char_p(binascii.hexlify(packet[self.rtp_header_size:]))

  def __str__(self):
    packet = 'RTP packet PT=%u SEQ=%u TS=%u SSRC=%x'
    packet = packet % (self.pt, self.seq, self.ts, self.ssrc)
    return packet


from ruamel.yaml import YAML, dump, RoundTripDumper
from raisimGymTorch.env.bin import test
from raisimGymTorch.env.RaisimGymVecEnv import RaisimGymVecEnv as VecEnv
import os
import math
import time
import raisimGymTorch.algo.ppo.module as ppo_module
import torch.nn as nn
import numpy as np
import torch
import argparse
import pdb
from raisimGymTorch.env.envs.train.model import Forward_Dynamics_Model
from raisimGymTorch.env.envs.train.action import UserCommand, Stochastic_action_planner_uniform_bin, Stochastic_action_planner_w_ITS
from raisimGymTorch.env.envs.train.model import Informed_Trajectory_Sampler_inference
from raisimGymTorch.env.envs.train.storage import Buffer
import random
from raisimGymTorch.env.envs.train.global_planner import Analytic_planner
from fastdtw import fastdtw
from dtw import dtw
from scipy.spatial.distance import euclidean
from joblib import Parallel, delayed

"""
Point-Goal Navigation using Time-correlated random sampler & Informed Trajectory Sampler (ITS)

- below code simply merge samples from two different sampler and optimize using MPPI
- set number of samples from Time-correlated random sampler to 0 if you only want to sample using ITS

"""

def transform_coordinate_LW(w_init_coordinate, l_coordinate_traj):
    """
    Transform LOCAL frame coordinate trajectory to WORLD frame coordinate trajectory
    (LOCAL frame --> WORLD frame)

    :param w_init_coordinate: initial coordinate in WORLD frame (1, coordinate_dim)
    :param l_coordinate_traj: coordintate trajectory in LOCAL frame (n_step, coordinate_dim)
    :return:
    """
    transition_matrix = np.array([[np.cos(w_init_coordinate[0, 2]), np.sin(w_init_coordinate[0, 2])],
                                  [- np.sin(w_init_coordinate[0, 2]), np.cos(w_init_coordinate[0, 2])]], dtype=np.float32)
    w_coordinate_traj = np.matmul(l_coordinate_traj, transition_matrix)
    w_coordinate_traj += w_init_coordinate[:, :-1]
    return w_coordinate_traj

def transform_coordinate_WL(w_init_coordinate, w_coordinate_traj):
    """
    Transform WORLD frame coordinate trajectory to LOCAL frame coordinate trajectory
    (WORLD frame --> LOCAL frame)

    :param w_init_coordinate: initial coordinate in WORLD frame (1, coordinate_dim)
    :param w_coordinate_traj: coordintate trajectory in WORLD frame (n_step, coordinate_dim)
    :return:
        l_coordinate_traj: coordintate trajectory in LOCAL frame (n_step, coordinate_dim)
    """
    transition_matrix = np.array([[np.cos(w_init_coordinate[0, 2]), np.sin(w_init_coordinate[0, 2])],
                                  [- np.sin(w_init_coordinate[0, 2]), np.cos(w_init_coordinate[0, 2])]], dtype=np.float32)
    l_coordinate_traj = w_coordinate_traj - w_init_coordinate[:, :-1]
    l_coordinate_traj = np.matmul(l_coordinate_traj, transition_matrix.T)
    return l_coordinate_traj

def compute_dtw(x, y, type="fast"):
    if type == "normal":
        dtw_distance = dtw(x, y, dist_method=euclidean).distance
        normalize_factor = (x.shape[0] + y.shape[0]) / 2
        return dtw_distance / normalize_factor
    elif type == "not_normal":
        dtw_distance = dtw(x, y, dist_method=euclidean).distance
        return dtw_distance
    else:
        return fastdtw(x, y, dist=euclidean)

# task specification
task_name = "point_goal_nav"

# configuration
parser = argparse.ArgumentParser()
parser.add_argument('-fw', '--fdm_weight', help='Forward Dynamics Model weight path', type=str, default='')
parser.add_argument('-iw', '--its_weight', help='Informed Trajectory Sampler weight path', type=str, default='')
parser.add_argument('-tw', '--tracking_weight', help='velocity command tracking policy weight path', type=str, default='')
parser.add_argument('-v', '--validate', help='validation or test', type=bool, default=False)
parser.add_argument('-cf', '--cf', help='extra note', type=str, default='')
args = parser.parse_args()
FDM_weight_path = args.fdm_weight
ITS_weight_path = args.its_weight
command_tracking_weight_path = args.tracking_weight
validation = args.validate
extra_note = args.cf

evaluate_seed = 10
random.seed(evaluate_seed)
np.random.seed(evaluate_seed)
torch.manual_seed(evaluate_seed)

# check if gpu is available
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# directories
task_path = os.path.dirname(os.path.realpath(__file__))
home_path = task_path + "/../../../.."

# config
cfg = YAML().load(open(task_path + "/cfg.yaml", 'r'))

# complete path configuration
if FDM_weight_path == '':
    assert cfg["path"]["home"] != '' or cfg["path"]["FDM"] != '', "Weight path configuration not complete."
    FDM_weight_path = cfg["path"]["home"] + cfg["path"]["FDM"]
if ITS_weight_path == '':
    assert cfg["path"]["home"] != '' or cfg["path"]["ITS"] != '', "Weight path configuration not complete."
    ITS_weight_path = cfg["path"]["home"] + cfg["path"]["ITS"]
if command_tracking_weight_path == '':
    assert cfg["path"]["home"] != '' or cfg["path"]["command_tracking"] != '', "Weight path configuration not complete."
    command_tracking_weight_path = cfg["path"]["home"] + cfg["path"]["command_tracking"]

# load model architecture
cfg["environment"]["test_initialize"]["point_goal"] = True
cfg["environment"]["test_initialize"]["safety_control"] = False

# user command sampling
user_command = UserCommand(cfg, cfg['CVAE_path_track']['planner']['wo_CVAE_number_of_sample'])

# create environment from the configuration file
cfg['environment']['num_envs'] = 1

# create environment from the configuration file
env = VecEnv(test.RaisimGymEnv(home_path + "/rsc", dump(cfg['environment'], Dumper=RoundTripDumper)), cfg['environment'], normalize_ob=False)

# shortcuts
user_command_dim = 3
proprioceptive_sensor_dim = 81
lidar_dim = 360
state_dim = cfg["environment_model"]["architecture"]["state_encoder"]["input"]
command_period_steps = math.floor(cfg['command_tracking']['command_period'] / cfg['environment']['control_dt'])
assert env.num_obs == proprioceptive_sensor_dim + lidar_dim, "Check configured sensor dimension"

# Use naive concatenation for encoding COM vel history
COM_feature_dim = cfg["environment_model"]["architecture"]["COM_encoder"]["naive"]["input"]
COM_history_time_step = cfg["environment_model"]["architecture"]["COM_encoder"]["naive"]["time_step"]
COM_history_update_period = int(cfg["environment_model"]["architecture"]["COM_encoder"]["naive"]["update_period"] / cfg["environment"]["control_dt"])
assert state_dim - lidar_dim == COM_feature_dim * COM_history_time_step, "Check COM_encoder output and state_encoder input in the cfg.yaml"

command_tracking_ob_dim = user_command_dim + proprioceptive_sensor_dim
command_tracking_act_dim = env.num_acts

COM_buffer = Buffer(env.num_envs, COM_history_time_step, COM_feature_dim)

# Load pre-trained command tracking policy weight
assert command_tracking_weight_path != '', "Pre-trained command tracking policy weight path should be determined."
command_tracking_policy = ppo_module.MLP(cfg['command_tracking']['architecture'], nn.LeakyReLU,
                                         command_tracking_ob_dim, command_tracking_act_dim)
command_tracking_policy.load_state_dict(torch.load(command_tracking_weight_path, map_location=device)['actor_architecture_state_dict'])
command_tracking_policy.to(device)
command_tracking_weight_dir = command_tracking_weight_path.rsplit('/', 1)[0] + '/'
iteration_number = command_tracking_weight_path.rsplit('/', 1)[1].split('_', 1)[1].rsplit('.', 1)[0]
env.load_scaling(command_tracking_weight_dir, int(iteration_number))

print("Loaded command tracking policy weight from {}\n".format(command_tracking_weight_path))

# Load learned environment model weight
loaded_environment_model = Forward_Dynamics_Model(state_encoding_config=cfg["environment_model"]["architecture"]["state_encoder"],
                                                  command_encoding_config=cfg["environment_model"]["architecture"]["command_encoder"],
                                                  recurrence_config=cfg["environment_model"]["architecture"]["recurrence"],
                                                  prediction_config=cfg["environment_model"]["architecture"]["traj_predictor"],
                                                  device=device)
loaded_environment_model.load_state_dict(torch.load(FDM_weight_path, map_location=device)['model_architecture_state_dict'])
loaded_environment_model.eval()
loaded_environment_model.to(device)
print("Loaded Forward Dynamics Model weight from {}\n".format(FDM_weight_path))

# Load sampler
n_prediction_step = int(cfg["CVAE_path_track"]["planner"]["prediction_period"] / cfg['command_tracking']['command_period'])
wo_cvae_sampler = Stochastic_action_planner_uniform_bin(command_range=cfg["environment"]["command"],
                                                        n_sample=cfg["CVAE_path_track"]["planner"]["wo_CVAE_number_of_sample"],
                                                        n_horizon=n_prediction_step,
                                                        n_bin=cfg["CVAE_path_track"]["planner"]["number_of_bin"],
                                                        beta=cfg["CVAE_path_track"]["planner"]["wo_CVAE_beta"],
                                                        gamma=cfg["CVAE_path_track"]["planner"]["gamma"],
                                                        sigma=cfg["CVAE_path_track"]["planner"]["sigma"],
                                                        noise_sigma=0.1,
                                                        noise=False,
                                                        action_dim=user_command_dim,
                                                        random_command_sampler=user_command)

w_cvae_sampler = Informed_Trajectory_Sampler_inference(
    latent_dim=cfg["CVAE_path_track"]["architecture"]["latent_dim"],
    state_encoding_config=cfg["CVAE_path_track"]["architecture"]["state_encoder"],
    waypoint_encoding_config=cfg["CVAE_path_track"]["architecture"]["waypoint_encoder"],
    waypoint_recurrence_encoding_config=cfg["CVAE_path_track"]["architecture"]["waypoint_recurrence_encoder"],
    latent_decoding_config=cfg["CVAE_path_track"]["architecture"]["latent_decoder"],
    recurrence_decoding_config=cfg["CVAE_path_track"]["architecture"]["recurrence_decoder"],
    command_decoding_config=cfg["CVAE_path_track"]["architecture"]["command_decoder"],
    device=device,
    trained_weight=ITS_weight_path,
    cfg_command=cfg["environment"]["command"]
)
w_cvae_sampler.eval()
w_cvae_sampler.to(device)
print("Loaded Informed Trajectory Sampler weight from {}\n".format(ITS_weight_path))

# Set action planner
action_planner = Stochastic_action_planner_w_ITS(wo_cvae_sampler=wo_cvae_sampler,
                                                 w_cvae_sampler=w_cvae_sampler,
                                                 wo_cvae_n_sample=cfg["CVAE_path_track"]["planner"]["wo_CVAE_number_of_sample"],
                                                 w_cvae_n_sample=cfg["CVAE_path_track"]["planner"]["CVAE_number_of_sample"],
                                                 n_prediction_step=n_prediction_step,
                                                 gamma=cfg["CVAE_path_track"]["planner"]["gamma"],
                                                 beta=cfg["CVAE_path_track"]["planner"]["CVAE_beta"])

# MUST safe period from collision
MUST_safety_period = 3.0
MUST_safety_period_n_steps = int(MUST_safety_period / cfg['command_tracking']['command_period'])

# Set constant
collision_threshold = 0.3
goal_distance_threshold = 10
num_goals = cfg["environment"]["n_goals_per_env"]
if validation:
    init_seed = cfg["environment"]["seed"]["validate"]
    print("Validating ...")
else:
    init_seed = cfg["environment"]["seed"]["evaluate"]
    print("Evaluating ...")
goal_time_limit = 180.

# hyperparameter for path tracking
lookahead_path_length = 4.8
nDTX_temperature = 8

# Make directory to save results
num_cvae_sample = cfg["CVAE_path_track"]["planner"]["CVAE_number_of_sample"]
num_wo_cvae_sample = cfg["CVAE_path_track"]["planner"]["wo_CVAE_number_of_sample"]

# Empty container to check time
time_check = []

print("<<-- Evaluating Point Goal Navigation -->>")

pdb.set_trace()

for grid_size in [2.3, 3., 4., 5.]:
    print("===========================================")
    print(f"Grid_{str(grid_size)}:")
    print("===========================================")

    # Set obstacle grid size
    cfg["environment"]["test_obstacle_grid_size"] = grid_size

    for env_id in range(cfg["environment"]["n_evaluate_envs"]):
        # Generate new environment with different seed (reset is automatically called)
        cfg["environment"]["seed"]["evaluate"] = env_id * 10 + init_seed
        env = VecEnv(test.RaisimGymEnv(home_path + "/rsc", dump(cfg['environment'], Dumper=RoundTripDumper)), cfg['environment'], normalize_ob=False)
        env.load_scaling(command_tracking_weight_dir, int(iteration_number))

        # set analytic planner
        map_size = np.zeros(2, dtype=np.float32)
        env.get_map_size(map_size)
        assert map_size[0] == map_size[1], "Map should be a square"
        analytic_planner = Analytic_planner(env, map_size[0], cfg["environment"]["analytic_planner"]["planning_time"], seed=evaluate_seed)

        n_test_case = 0
        n_path_failure = 0  # Count number of finding path failure

        while n_test_case < num_goals:
            # Reset
            env.initialize_n_step()
            env.reset()
            action_planner.reset()
            goal_position = env.set_goal()[np.newaxis, :]
            COM_buffer.reset()

            # Initialize
            step = 0
            sample_user_command = np.zeros(3)
            goal_current_duration = 0.
            traversal_distance = 0.
            previous_coordinate = None
            current_coordinate = None

            # plan path with analytic planner
            start_position = env.coordinate_observe()[0, :2]
            planned_path = analytic_planner.plan(start_position.astype(np.float64), goal_position[0, :].astype(np.float64))

            if cfg["environment"]["render"] and cfg["environment"]["analytic_planner"]["visualize"] and planned_path is not None:
                analytic_planner.visualize_path()

            if planned_path is None:
                n_path_failure += 1
                if n_path_failure > 10:
                    break
            else:
                # Available test case found
                n_test_case += 1

                # Save COM trajectory to compute path similarity with DTW
                traversed_path = []

                while True:
                    frame_start = time.time()
                    control_start = time.time()

                    new_action_time = step % command_period_steps == 0

                    # log traversal distance (just env 0)
                    previous_coordinate = current_coordinate
                    current_coordinate = env.coordinate_observe()
                    if previous_coordinate is not None:
                        delta_coordinate = current_coordinate[0, :-1] - previous_coordinate[0, :-1]
                        traversal_distance += (delta_coordinate[0] ** 2 + delta_coordinate[1] ** 2) ** 0.5

                    # observation before taking step
                    obs, _ = env.observe(False)

                    # update COM feature
                    if step % COM_history_update_period == 0:
                        COM_feature = np.concatenate((obs[:, :3], obs[:, 15:21]), axis=1)
                        COM_buffer.update(COM_feature)

                    if new_action_time:
                        if step != 0:
                            # Pass the one-step planned command to the low-level controller
                            sample_user_command = sample_user_command_traj[1, :]

                        # plan_start = time.time()

                        # prepare state
                        init_coordinate_obs = env.coordinate_observe()
                        lidar_data = obs[0, proprioceptive_sensor_dim:]
                        COM_history_feature = COM_buffer.return_data(flatten=True)[0, :]
                        state = np.concatenate((lidar_data, COM_history_feature)).astype(np.float32)

                        # compute final goal position w/ respect to local frame
                        goal_position_L = transform_coordinate_WL(init_coordinate_obs, goal_position)
                        current_goal_distance = np.sqrt(np.sum(np.power(goal_position_L, 2)))
                        goal_position_L *= np.clip(goal_distance_threshold / current_goal_distance, a_min=None, a_max=1.)
                        goal_position_L = goal_position_L.astype(np.float32)
                        current_goal_distance = np.sqrt(np.sum(np.power(goal_position_L, 2)))

                        # use path similarity as reward for path tracking
                        waypoints_list_L = transform_coordinate_WL(init_coordinate_obs, planned_path)
                        waypoint_distance_list = np.sqrt(np.sum(np.power(waypoints_list_L, 2), axis=-1))
                        closest_waypoint_idx = np.argmin(waypoint_distance_list)
                        waypoints_list_L_closest = waypoints_list_L - waypoints_list_L[closest_waypoint_idx]
                        along_waypoint_distance = 0
                        farthest_waypoint_idx = planned_path.shape[0] - 1
                        for idx in range(closest_waypoint_idx, waypoints_list_L_closest.shape[0] - 1):
                            along_waypoint_distance += np.linalg.norm(waypoints_list_L_closest[idx + 1] - waypoints_list_L_closest[idx], ord=2)
                            if along_waypoint_distance > lookahead_path_length:
                                farthest_waypoint_idx = idx + 1
                                break

                        waypoint_idx = list(range(closest_waypoint_idx, farthest_waypoint_idx + 1))
                        waypoints_list_L = waypoints_list_L[waypoint_idx]  # (n_waypoint, 2)

                        if cfg["environment"]["render"] and cfg["environment"]["analytic_planner"]["visualize"]:
                            env.visualize_waypoints(planned_path[[closest_waypoint_idx, farthest_waypoint_idx]])

                        # interpolate for sparsely sampled points
                        if farthest_waypoint_idx - closest_waypoint_idx < 4 and along_waypoint_distance > lookahead_path_length:
                            if farthest_waypoint_idx - closest_waypoint_idx == 3:
                                interpolated_waypoints = []
                                for idx in range(farthest_waypoint_idx - closest_waypoint_idx):
                                    interpolated_waypoints.append(np.linspace(waypoints_list_L[idx], waypoints_list_L[idx + 1], num=3)[:-1, :])
                                interpolated_waypoints.append(waypoints_list_L[-1, :][np.newaxis, :])
                                waypoints_list_L = np.concatenate(interpolated_waypoints, axis=0)
                            elif farthest_waypoint_idx - closest_waypoint_idx == 2:
                                interpolated_waypoints = []
                                for idx in range(farthest_waypoint_idx - closest_waypoint_idx):
                                    interpolated_waypoints.append(np.linspace(waypoints_list_L[idx], waypoints_list_L[idx + 1], num=4)[:-1, :])
                                interpolated_waypoints.append(waypoints_list_L[-1, :][np.newaxis, :])
                                waypoints_list_L = np.concatenate(interpolated_waypoints, axis=0)
                            elif farthest_waypoint_idx - closest_waypoint_idx == 1:
                                waypoints_list_L = np.linspace(waypoints_list_L[0], waypoints_list_L[1], num=8)

                        # sample command sequences
                        # one-step lookahead planning (first step command is already determined)
                        action_candidates = action_planner.sample(observation=torch.from_numpy(state).unsqueeze(0).to(device),
                                                                  goal_position=None,
                                                                  waypoints=torch.from_numpy(waypoints_list_L[:, np.newaxis, :]).to(device),
                                                                  current_command=None)
                        action_candidates[0, :, :] = sample_user_command
                        action_planner.sampled_command_traj = action_candidates

                        # simulate sampled command sequences
                        state = np.tile(state, (num_wo_cvae_sample + num_cvae_sample, 1))
                        predicted_P_cols, predicted_coordinates = loaded_environment_model(torch.from_numpy(state).to(device),
                                                                                torch.from_numpy(action_candidates).to(device),
                                                                                training=False)
                        predicted_P_cols = np.squeeze(predicted_P_cols, axis=-1)

                        # Hard constraint for collision
                        for sample_id in range(num_wo_cvae_sample + num_cvae_sample):
                            current_done = np.where(predicted_P_cols[:, sample_id] > collision_threshold)[0]
                            if len(current_done) != 0:
                                done_idx = np.min(current_done)
                                predicted_coordinates[done_idx + 1:, sample_id, :] = predicted_coordinates[done_idx, sample_id, :]
                                predicted_P_cols[done_idx + 1:, sample_id] = predicted_P_cols[done_idx, sample_id]

                        # compute reward (goal reward + safety reward)
                        if farthest_waypoint_idx - closest_waypoint_idx < 4 and along_waypoint_distance < lookahead_path_length:
                            # if the robot is very near the goal, give reward directly with goal position
                            delta_goal_distance = current_goal_distance - np.sqrt(np.sum(np.power(predicted_coordinates - goal_position_L, 2), axis=-1))
                            goal_reward = np.sum(delta_goal_distance, axis=0)
                        else:
                            padded_predicted_coordinates = np.concatenate((np.zeros((1, predicted_coordinates.shape[1], predicted_coordinates.shape[2])), predicted_coordinates), axis=0)
                            results = Parallel(n_jobs=8)(delayed(compute_dtw)(waypoints_list_L, padded_predicted_coordinates[:, i, :]) for i in range(num_wo_cvae_sample + num_cvae_sample))  # takes about 0.1 [s]
                            dtw_distance = np.array([r[0] for r in results])
                            dtw_distance -= np.min(dtw_distance)
                            goal_reward = np.exp(- dtw_distance / nDTX_temperature)
                            goal_reward = np.array(goal_reward)

                        goal_reward -= np.min(goal_reward)
                        goal_reward /= (np.max(goal_reward) + 1e-5)  # normalize reward

                        safety_reward = 1 - predicted_P_cols
                        safety_reward = np.mean(safety_reward, axis=0)
                        safety_reward /= (np.max(safety_reward) + 1e-5)  # normalize reward

                        reward = 1.0 * goal_reward + 1.0 * safety_reward

                        # exclude trajectory that collides with obstacle
                        coll_idx = np.where(np.sum(np.where(predicted_P_cols[:MUST_safety_period_n_steps, :] > collision_threshold, 1, 0), axis=0) != 0)[0]
                        if len(coll_idx) != (num_wo_cvae_sample + num_cvae_sample):
                            reward[coll_idx] = 0

                        # optimize command sequence
                        _, sample_user_command_traj = action_planner.action(reward)

                        # simulate optimized command sequence
                        state = state[0, :][np.newaxis, :]
                        predicted_P_cols, predicted_coordinates = loaded_environment_model(torch.from_numpy(state).to(device),
                                                                                torch.from_numpy(sample_user_command_traj[:, np.newaxis, :]).to(device),
                                                                                training=False)

                        # Hard constraint for collision
                        predicted_P_cols = np.squeeze(predicted_P_cols, axis=-1)
                        current_done = np.where(predicted_P_cols[:, 0] > collision_threshold)[0]
                        if len(current_done) != 0:
                            done_idx = np.min(current_done)
                            predicted_coordinates[done_idx + 1:, 0, :] = predicted_coordinates[done_idx, 0, :]
                            predicted_P_cols[done_idx + 1:, 0] = predicted_P_cols[done_idx, 0]

                        # visualize predicted modified command trajectory
                        if cfg["environment"]["render"]:
                            w_coordinate_modified_command_path = transform_coordinate_LW(init_coordinate_obs, predicted_coordinates[:, 0, :])
                            P_col_modified_command_path = predicted_P_cols[:, 0][:, np.newaxis]
                            env.visualize_modified_command_traj(w_coordinate_modified_command_path,
                                                                P_col_modified_command_path,
                                                                collision_threshold)

                        # Save COM coordinate
                        traversed_path.append(init_coordinate_obs[0, :-1])  # (coordinate_dim,)

                        # plan_end = time.time()

                        # print(plan_end - plan_start)

                    # Execute first command in optimized command sequence using command tracking controller
                    tracking_obs = np.concatenate((sample_user_command, obs[0, :proprioceptive_sensor_dim]))[np.newaxis, :]
                    tracking_obs = env.force_normalize_observation(tracking_obs, type=1)
                    tracking_obs = tracking_obs.astype(np.float32)

                    with torch.no_grad():
                        tracking_action = command_tracking_policy.architecture(torch.from_numpy(tracking_obs).to(device))

                    control_end = time.time()

                    _, done = env.step(tracking_action.cpu().detach().numpy())

                    # Update progress
                    step += 1
                    goal_current_duration += cfg['environment']['control_dt']

                    frame_end = time.time()

                    # # Check time
                    # if new_action_time:
                    #     time_check.append(control_end - control_start)
                    #     if len(time_check) == 5000:
                    #         time_check = np.array(time_check)
                    #         print(f"Mean: {np.mean(time_check[50:])}")
                    #         print(f"Std: {np.std(time_check[50:])}")
                    #         pdb.set_trace()

                    # if new_action_time:
                    #    print(control_end - control_start)

                    # if new_action_time:
                    #     print(frame_end - frame_start)

                    if cfg["realistic"]:
                        wait_time = cfg['environment']['control_dt'] - (frame_end-frame_start)
                        if wait_time > 0.:
                            time.sleep(wait_time)

                    if goal_current_duration > goal_time_limit:
                        done[0] = True

                    # fail
                    if done[0] == True:
                        break
                    # success
                    elif current_goal_distance < 0.6:
                        if cfg["environment"]["render"] and cfg["environment"]["visualize_path"]:
                            pdb.set_trace()

                        # Update
                        traversed_path = np.stack(traversed_path, axis=0)
                        total_normal_dtw = compute_dtw(planned_path, traversed_path, type="normal")

                        print(f"Time: {round(step * cfg['environment']['control_dt'], 1)}")
                        print(f"Distance: {round(traversal_distance, 1)}")
                        print(f"Normal DTW: {round(total_normal_dtw, 1)}")
                        print("-------------------------------")

                        break

# env.stop_video_recording()
env.turn_off_visualization()


# 383. Ransom Note
# Given two strings ransomNote and magazine, return true if ransomNote can be constructed by using the letters from magazine and false otherwise.
# Each letter in magazine can only be used once in ransomNote.

class Solution:
    def canConstruct(self, ransomNote: str, magazine: str) -> bool:
        record: list[int] = [0] * 26
        for i in magazine:
            record[ord(i) - ord('a')] += 1
        for i in ransomNote:
            record[ord(i) - ord('a')] -= 1
        for i in range(26):
            if record[i] < 0:
                return False
        return True


from __future__ import absolute_import

import re

from touchdown.core import argument, errors, resource, serializers

from . import provisioner

try:
    import fuselage
    from fuselage import argument as f_args, builder, bundle, resources
except ImportError:
    raise errors.Error(
        "You need the fuselage package to use the fuselage_bundle resource"
    )

def underscore(title):
    return re.sub(r"(?<=[a-z])(?=[A-Z])", u"_", title).lower()

arguments = {
    f_args.Boolean: lambda resource_type, klass, arg: argument.Boolean(field=arg),
    f_args.String: lambda resource_type, klass, arg: argument.String(field=arg),
    f_args.FullPath: lambda resource_type, klass, arg: argument.String(field=arg),
    f_args.File: lambda resource_type, klass, arg: argument.String(field=arg),
    f_args.Integer: lambda resource_type, klass, arg: argument.Integer(field=arg),
    f_args.Octal: lambda resource_type, klass, arg: argument.Integer(field=arg),
    f_args.Dict: lambda resource_type, klass, arg: argument.Dict(field=arg),
    f_args.List: lambda resource_type, klass, arg: argument.List(field=arg),
    f_args.SubscriptionArgument: lambda resource_type, klass, arg: argument.List(
        field=arg
    ),
    f_args.PolicyArgument: lambda resource_type, klass, arg: argument.String(
        field=arg, choices=resource_type.policies.keys()
    ),
}

class FuselageResource(resource.Resource):
    @classmethod
    def adapt(base_klass, resource_type):
        args = {
            "resource_name": underscore(resource_type.__resource_name__),
            "fuselage_class": resource_type,
            "root": argument.Resource(Bundle),
        }

        for arg, klass in resource_type.__args__.items():
            args[arg] = arguments[klass.__class__](resource_type, klass, arg)

        cls = type(resource_type.__resource_name__, (base_klass,), args)

        def _(self, **kwargs):
            arguments = {"parent": self}
            arguments.update(kwargs)
            resource = cls(**arguments)
            if not self.resources:
                self.resources = []
            self.resources.append(resource)
            self.add_dependency(resource)
            return resource

        setattr(Bundle, "add_%s" % cls.resource_name, _)

        return cls

class BundleSerializer(serializers.Serializer):
    def render(self, runner, value):
        b = bundle.ResourceBundle()
        for res in value:
            b.add(res.fuselage_class(**serializers.Resource().render(runner, res)))
        return builder.build(b)

    def pending(self, runner, value):
        for res in value:
            if serializers.Resource().pending(runner, res):
                return True
        return False

class Bundle(provisioner.Provisioner):

    resource_name = "fuselage_bundle"

    always_apply = argument.Boolean()
    resources = argument.List(
        argument.Resource(FuselageResource),
        field="script",
        serializer=BundleSerializer(),
    )
    sudo = argument.Boolean(field="sudo", default=True)

class Describe(provisioner.Describe):

    name = "describe"
    resource = Bundle

    def describe_object(self):
        if self.resource.always_apply:
            return {"Results": "Pending"}

        if not self.resource.target:
            # If target is not set we are probably dealing with an AMI... YUCK
            # Bail out
            return {"Result": "Pending"}

        serializer = serializers.Resource()
        if serializer.pending(self.runner, self.resource):
            return {"Result": "Pending"}

        kwargs = serializer.render(self.runner, self.resource)

        try:
            client = self.runner.get_plan(self.resource.target).get_client()
        except errors.ServiceNotReady:
            return {"Result": "Pending"}

        try:
            client.run_script(kwargs["script"], ["-s"])
        except errors.CommandFailed as e:
            if e.exit_code == 254:
                return {"Result": "Success"}

        return {"Result": "Pending"}

class Apply(provisioner.Apply):

    resource = Bundle

for attr, value in vars(resources).items():
    if type(value) == fuselage.resource.ResourceType:
        locals()[attr] = FuselageResource.adapt(value)


import argparse
import math
import os
import sys

import matplotlib.pyplot as plt

from typing import Dict, List, Tuple

_default_timings = "read,tracking,integration,raycast,render"
_default_plots = "timings"
_iteration_column_name = "iteration"

class QuantityValues:
    """Contains the values of some quantity for a single run."""
    def __init__(self, name: str, unit: str) -> None:
        self.name = name
        self.unit = unit
        self.values = []

    def __str__(self) -> str:
        if self.unit:
            return "{} ({})".format(self.name, self.unit)
        else:
            return "{}".format(self.name)

    def append(self, value: float) -> None:
        self.values.append(value)

    def mean(self) -> float:
        return sum(self.values) / len(self.values)

class Stats:
    """Statistics for a single run."""
    def __init__(self, filename: str) -> None:
        self.filename = filename
        # Open a file or stdin
        if not self.filename or self.filename == "-":
            f = sys.stdin
        else:
            f = open(filename)
        # Read the header and initialize the dictionaries.
        header = f.readline()
        self.stats_by_name, self.stats_by_idx = Stats._init_stats(header)
        # Read the data for each iteration.
        for line in f:
            Stats._append_stats(line, self.stats_by_idx)
        if f is not sys.stdin:
            f.close()

    def max_time(self, timings: List[str]) -> List[float]:
        # Combine the values of all specified timings in a list of lists.
        timing_values = [self.stats_by_name[t].values for t in timings]
        # Find the maximum time per iteration using the splat (*) operator.
        return [max(x) for x in zip(*timing_values)]

    def total_time(self, timings: List[str]) -> List[float]:
        # Combine the values of all specified timings in a list of lists.
        timing_values = [self.stats_by_name[t].values for t in timings]
        # Sum the timing values to get the total frame time per iteration using
        # the splat (*) operator.
        return [sum(x) for x in zip(*timing_values)]

    @staticmethod
    def _init_stats(header: str) -> Tuple[Dict[str,QuantityValues],Dict[int,QuantityValues]]:
        """Given the TSV header return two dictionaries to access the individual
        quantities by name or zero-based column index."""
        stats_by_name = {}
        stats_by_idx = {}
        columns = header.strip().split("\t")
        for i, c in enumerate(columns):
            # Extract the name and its unit.
            unit_idx = c.rfind(" (")
            if unit_idx >= 0:
                name = c[0:unit_idx].strip()
                unit = c[unit_idx+2:-1].strip()
            else:
                name = c.strip()
                unit = ""
            # Set the unit to empty if it's not a unit.
            if unit in ["#", "-", "?"]:
                unit = ""
            stats_by_name[name] = QuantityValues(name, unit)
            # The following will create a reference, not a copy.
            stats_by_idx[i] = stats_by_name[name]
        return stats_by_name, stats_by_idx

    @staticmethod
    def _append_stats(line: str, stats_by_idx: Dict[int,QuantityValues]) -> None:
        columns = line.strip().split("\t")
        for i, c in enumerate(columns):
            if c != "*":
                stats_by_idx[i].append(float(c))

class Plotter:
    """Creates plots from statistics for multiple runs."""
    plot_names = ["timings", "path"]
    plot_names_str = ", ".join(plot_names)

    def __init__(self, stats: List[Stats], args) -> None:
        # Create dictionaries from plot names to class methods.
        self._plots_to_plot_funcs = dict(zip(Plotter.plot_names,
            [self._plot_timings, self._plot_path]))
        self._plots_to_limit_funcs = dict(zip(Plotter.plot_names,
            [self._limits_timings, self._limits_path]))

        self.stats = stats
        self.plots = args.plots
        self.timings = args.timings
        self.lineplot = args.lineplot
        self.save = args.save
        self.quiet = args.quiet

        # Compute the axis limits for all plots.
        self._limits = {plot_name: self._plots_to_limit_funcs[plot_name]() for plot_name in self.plots}

        # Create a figure for each run.
        for s in self.stats:
            self._plot(s)
            if self.save:
                # Generate the filename of the PNG file.
                if not s.filename or s.filename == "-":
                    filename_png = "figure.png"
                else:
                    filename_png = s.filename + ".png"
                figure = plt.gcf()
                figure.set_size_inches(16, 9)
                plt.savefig(filename_png, dpi = 200, bbox_inches="tight")
        if not self.quiet:
            plt.show()

    def _limits_timings(self) -> List[float]:
        x_min = min([s.stats_by_name[_iteration_column_name].values[0] for s in self.stats])
        x_max = max([s.stats_by_name[_iteration_column_name].values[-1] for s in self.stats])
        if self.lineplot:
            y_max = max([max(s.max_time(self.timings)) for s in self.stats])
        else:
            y_max = max([max(s.total_time(self.timings)) for s in self.stats])
        return [x_min, x_max, 0, y_max]

    def _limits_path(self) -> List[float]:
        x_min = min([min(s.stats_by_name["tx"].values) for s in self.stats])
        x_max = max([max(s.stats_by_name["tx"].values) for s in self.stats])
        y_min = min([min(s.stats_by_name["ty"].values) for s in self.stats])
        y_max = max([max(s.stats_by_name["ty"].values) for s in self.stats])
        return [x_min, x_max, y_min, y_max]

    def _plot_timings(self, stats: Stats, axis) -> None:
        # ASSUMPTION: the time units are seconds.
        axis.set_title("Computation time")
        axis.set_xlabel("Frame")
        axis.set_ylabel("Time (ms)")
        iterations = stats.stats_by_name[_iteration_column_name].values
        # Stack all the timing data together.
        timing_data = [[1000 * x for x in stats.stats_by_name[t].values] for t in self.timings]
        if self.lineplot:
            for i, t in enumerate(self.timings):
                axis.plot(iterations, timing_data[i], label=t)
        else:
            axis.stackplot(iterations, timing_data, labels=self.timings)
        axis.grid(True)
        axis.legend(loc="upper left")
        l = self._limits["timings"]
        axis.set_xlim(l[0], l[1])
        axis.set_ylim(l[2], 1000 * l[3])

    def _plot_path(self, stats: Stats, axis) -> None:
        axis.set_title("Path")
        axis.set_xlabel("x (m)")
        axis.set_ylabel("y (m)")
        axis.plot(stats.stats_by_name["tx"].values,
                stats.stats_by_name["ty"].values)
        axis.plot(stats.stats_by_name["tx"].values[0],
                stats.stats_by_name["ty"].values[0], color="blue", marker="o",
                linestyle="", label="Start")
        axis.plot(stats.stats_by_name["tx"].values[-1],
                stats.stats_by_name["ty"].values[-1], color="blue", marker="x",
                linestyle="", label="End")
        axis.grid(True)
        axis.legend(loc="upper left")
        l = self._limits["path"]
        axis.set_xlim(l[0], l[1])
        axis.set_ylim(l[2], l[3])

    def _plot(self, stats: Stats) -> None:
        fig = plt.figure()
        fig.canvas.set_window_title(os.path.basename(stats.filename))

        num_plots = len(self.plots)
        for i, plot_name in enumerate(self.plots):
            ax = fig.add_subplot(num_plots, 1, i + 1)
            self._plots_to_plot_funcs[plot_name](stats, ax)

def parse_arguments():
    parser = argparse.ArgumentParser(description=(""))
    parser.add_argument("files", nargs="*", metavar="FILE", default=["-"],
            help=("""A TSV file containing the supereight log. With no FILE or
            when FILE is -, read standard input. If multiple files are provided
            each one will be plotted in a different window using the same axis
            limits for easy comparison."""))
    parser.add_argument("-t", "--timings", metavar="NAMES",
            default=_default_timings,
            help="""select which timings to show. NAMES is a comma-separated
            list containing any TSV column name (default: """ + _default_timings
            + ")")
    parser.add_argument("-p", "--plots", metavar="NAMES",
            default=_default_plots,
            help="""select which plots to show. NAMES is a comma-separated list
            which can contain any of the names: """ + Plotter.plot_names_str + """
            (default: """ + _default_plots + ")")
    parser.add_argument("-l", "--lineplot", action="store_true",
            help="""create a line plot instead of a stacked plot for timings""")
    parser.add_argument("-s", "--save", action="store_true",
            help="""save the plots using the same filenames appended with the
            .png extension""")
    parser.add_argument("-q", "--quiet", action="store_true",
            help="""hide all plots""")
    args = parser.parse_args()
    # Extract the timing names.
    args.timings = args.timings.split(",")
    # Extract the plot names and ensure they are valid.
    args.plots = args.plots.split(",")
    for plot_name in args.plots:
        if plot_name not in Plotter.plot_names:
            print("Error, invalid plot name: " + plot_name)
            print("Expected one of: " + Plotter.plot_names_str)
            sys.exit(2)
    return args

if __name__ == "__main__":
    try:
        args = parse_arguments()
        stats = [Stats(filename) for filename in args.files]
        plotter = Plotter(stats, args)
    except (KeyboardInterrupt, BrokenPipeError):
        pass


#importando bibliotecas
import pandas as pd
import plotly.express as px
import numpy as np
from haversine import haversine
import folium
import streamlit as st
from datetime import datetime
from PIL import Image
import folium
from folium import *
from streamlit_folium import folium_static
import sys
from utils.clean import Clean
from utils.sidebar import Side_Bar

# Definir cor de fundo
st.set_page_config(page_title='VisÃ£o Entregadores'
                   ,page_icon="ðŸšš"
                   ,layout="wide")

#=========
# FunÃ§Ãµes
#=========
def df_ratings(df1):
    """
    Calcula as mÃ©dias das avaliaÃ§Ãµes dos entregadores e retorna um DataFrame com os resultados.

    ParÃ¢metros
    ----------
    df1 : pandas.DataFrame
        DataFrame contendo os dados dos pedidos.

    Retorna
    -------
    pandas.DataFrame
        DataFrame com as mÃ©dias das avaliaÃ§Ãµes dos entregadores.

    """
    df_ratings =(
        df1.loc[: , ['Delivery_person_ID','Delivery_person_Ratings']]
                    .groupby('Delivery_person_ID')
                    .mean()
                    .reset_index()
                    )
    return df_ratings

def ratings(df1,columns,groupby,agg,kpis,new_columns):
    """
    Calcula as mÃ©tricas de avaliaÃ§Ã£o com base nos dados fornecidos e retorna um DataFrame com os resultados.

    ParÃ¢metros
    ----------
    df1 : pandas.DataFrame
        DataFrame contendo os dados dos pedidos.
    columns : list
        Lista das colunas a serem consideradas para o cÃ¡lculo das mÃ©tricas.
    groupby : str
        Coluna a ser usada para agrupar os dados.
    agg : str or dict
        FunÃ§Ã£o ou dicionÃ¡rio de funÃ§Ãµes de agregaÃ§Ã£o a serem aplicadas aos dados agrupados.
    kpis : list
        Lista de mÃ©tricas a serem calculadas.
    new_columns : list
        Lista dos novos nomes das colunas no DataFrame resultante.

    Retorna
    -------
    pandas.DataFrame
        DataFrame com as mÃ©tricas de avaliaÃ§Ã£o calculadas.

    """
    df_ratings = (
        df1.loc[ : , columns]
                    .groupby(groupby)
                    .agg({agg:kpis})
                    )
    df_ratings.columns = new_columns
    df_ratings = df_ratings.reset_index()
    return df_ratings

def df_top_fast_slow(df1,top_asc):
    """
    Retorna um DataFrame com os 10 principais entregadores mais rÃ¡pidos ou mais lentos em cada tipo de cidade.

    ParÃ¢metros
    ----------
    df1 : pandas.DataFrame
        DataFrame contendo os dados dos pedidos.
    top_asc : bool
        Indicador para classificar os resultados em ordem ascendente (True) ou descendente (False).

    Retorna
    -------
    pandas.DataFrame
        DataFrame com os 10 principais entregadores mais rÃ¡pidos ou mais lentos em cada tipo de cidade.

    """
    df2 = (
        df1.loc[:, ['Delivery_person_ID','City','Time_taken(min)']]
                    .groupby(['Delivery_person_ID','City'])
                    .min()
                    .sort_values(['City','Time_taken(min)'], ascending=top_asc)
                    .reset_index()
            )

    dfa1 =df2.loc[df2['City'] == 'Metropolitian', :].head(10)
    dfa2 = df2.loc[df2['City'] == 'Urban', :].head(10)
    dfa3 = df2.loc[df2['City'] == 'Semi-Urban', :].head(10)

    df = pd.concat([dfa1,dfa2,dfa3]).reset_index(drop=True)
    return df

#=========
# Extract
#=========

df = pd.read_csv('./dataset/train.csv',sep=',')

#============
# Transnform
#============

clean = Clean()
df1 = clean.clean_code(df)

#=========================
# Barra Lateral Streamlit
#=========================

st.header('Marketplace - VisÃ£o dos Entregadores')
image_path = './images/Logo.png'
image = Image.open( image_path )
st.sidebar.image(image,width=120)

#Chamando a classe e criando filtros na barra lateral
sidebar = Side_Bar()
date_slider, trafic_option, city_options, wheater_options = sidebar.sidebar(df1)

st.sidebar.markdown('___')
st.sidebar.markdown('## Powered by Thiago Vale')

# Ajustando o formato do date_slider
date_slider = date_slider.strftime('%Y-%m-%d')

# Filtrando o DataFrame
linhas_selecionadas = df1['Order_Date'] < date_slider
df1 = df1.loc[linhas_selecionadas, :]

linhas_selecionada = df1['Road_traffic_density'].isin(trafic_option)
df1 = df1.loc[linhas_selecionada,:]

linhas_selecionada = df1['City'].isin(city_options)
df1 = df1.loc[linhas_selecionada,:]

linhas_selecionada = df1['Weatherconditions'].isin(wheater_options)
df1 = df1.loc[linhas_selecionada,:]

#==================
# Layout Streamlit
#==================

tab1,tab2,tab3 = st.tabs(['VisÃ£o Gerencial',' ',' '])

with tab1:
    with st.container():
        st.title('Overall Metrics')

        col1,col2,col3,col4 = st.columns(4,gap='large')
        with col1:
            maior_idade =  df1.loc[:,'Delivery_person_Age'].max()
            col1.metric('Maior Idade',maior_idade)

        with col2:
            menor_idade = df1.loc[:,'Delivery_person_Age'].min()
            col2.metric('Menor Idade',menor_idade)

        with col3:
            melhor_condicao = df1.loc[:,'Vehicle_condition'].max()
            col3.metric('Melhor CondiÃ§Ã£o de Veiculos',melhor_condicao)

        with col4:
            pior_condicao = df1.loc[:,'Vehicle_condition'].min()
            col4.metric('Pior CondiÃ§Ã£o de Veiculos',pior_condicao)

    with st.container():
        st.markdown("""---""")
        st.title('AvaliaÃ§Ãµes')

        col1,col2 = st.columns( 2 )

        with col1:
            st.markdown('##### AvaliaÃ§Ã£o Media Por Entregador')
            st.dataframe(df_ratings(df1))

        with col2:
            st.markdown('##### AvaliaÃ§Ã£o Media Por Transito')
            df_ratings = (ratings(df1
                                      ,['Delivery_person_Ratings','Road_traffic_density']
                                      ,'Road_traffic_density'
                                      ,'Delivery_person_Ratings'
                                      ,['mean','std']
                                      ,['delivery_mean','delivery_std'])
                              )
            st.dataframe(df_ratings)

            st.markdown('##### AvaliaÃ§Ã£o Media Por Cilma')
            df_ratings = (ratings(df1
                                  ,['Delivery_person_Ratings','Weatherconditions']
                                  ,'Weatherconditions'
                                  ,'Delivery_person_Ratings'
                                  ,['mean','std']
                                  ,['delivery_mean','delivery_std'])
                          )
            st.dataframe(df_ratings)

    with st.container():
        st.markdown("""---""")
        st.title("Velocidade de Entrega")

        col1,col2 = st.columns( 2 )

        with col1:
            st.markdown('##### Top Entregadores mais rapidos')
            df_top = df_top_fast_slow(df1,True)
            st.dataframe(df_top.head(10))

        with col2:
            st.markdown('##### Top Entregadores mais lentos')

            df_tail = df_top_fast_slow(df1,False)
            st.dataframe(df_tail)


import unittest
import krpctest

class TestPartsFairing(krpctest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.new_save()
        cls.launch_vessel_from_vab('PartsFairing')
        cls.remove_other_vessels()
        parts = cls.connect().space_center.active_vessel.parts
        cls.fairing = parts.with_title(
            'AE-FF1 Airstream Protective Shell (1.25m)')[0].fairing

    def test_jettison(self):
        self.assertFalse(self.fairing.jettisoned)
        self.fairing.jettison()
        self.assertTrue(self.fairing.jettisoned)

if __name__ == '__main__':
    unittest.main()


import numpy as np
from random import shuffle

def get_divisors(n):
    return {div for div in range(2, n // 2 + 1) if n % div == 0} | {1, n}
    # return [1] + [div for div in range(2, n // 2 + 1) if n % div == 0] + [n]

def get_presents(n):
    return sum(get_divisors(n)) * 10

# Part 1

target = np.uint32(3_310_000)
start = 714_468
end = 840_840
presents = np.zeros(end-start+1, dtype="uint32")
houses = np.arange(start, end+1, dtype="uint32")

for elf in np.arange(1, end+1, dtype="uint32"):
    visits = (houses % elf == 0)
    presents[visits] += elf
    max_presents = presents.max()
    if (max_presents >= target):
        break
    print(f"{elf}: {target - max_presents:<7}", end="\r")

print()
number = houses[(presents >= target)]
print(f"Part 1: {number[0]}")   # Result: 776160

# start = 714_468
# end = 840_840 #850_080 #884_520 #1_000_000
# window = 100
# search_space = set(range(start, end))
# min_house = end

# try:
#     while True:
#         space = list(search_space)
#         shuffle(space)

#         houses = []
#         offset = 0
#         while not houses:
#             for i in range(window):
#                 if i == len(space): break
#                 house_number = space[offset+i]
#                 presents = get_presents(space[offset+i])

#                 if presents >= 33_100_000:
#                     houses.append(house_number)
#                 print(f"({offset+i}) House {house_number:<7}: {presents:<8} presents", end="\r")

#                 search_space.remove(house_number)

#             offset += window
#             if offset >= len(space): break

#         min_house = min(houses)
#         search_space -= set(range(min_house, end))
#         end = min_house
#         print("\n", houses, len(search_space))

#         if len(search_space) <= window: break

# except KeyboardInterrupt:
#     pass

# print(f"\nPart 1: {min_house}")