import math
import functools
from operator import mul

import torch
import torch.nn.functional as F
from torch import nn, einsum

from einops import rearrange, repeat, pack, unpack
from einops.layers.torch import Rearrange

from make_a_video_pytorch.attend import Attend

# helper functions

def exists(val):
    return val is not None

def default(val, d):
    return val if exists(val) else d

def mul_reduce(tup):
    return functools.reduce(mul, tup)

def divisible_by(numer, denom):
    return (numer % denom) == 0

mlist = nn.ModuleList

# for time conditioning

class SinusoidalPosEmb(nn.Module):
    def __init__(self, dim, theta = 10000):
        super().__init__()
        self.theta = theta
        self.dim = dim

    def forward(self, x):
        dtype, device = x.dtype, x.device
        assert dtype == torch.float, 'input to sinusoidal pos emb must be a float type'

        half_dim = self.dim // 2
        emb = math.log(self.theta) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, device = device, dtype = dtype) * -emb)
        emb = rearrange(x, 'i -> i 1') * rearrange(emb, 'j -> 1 j')
        return torch.cat((emb.sin(), emb.cos()), dim = -1).type(dtype)

# layernorm 3d

class RMSNorm(nn.Module):
    def __init__(self, chan, dim = 1):
        super().__init__()
        self.dim = dim
        self.gamma = nn.Parameter(torch.ones(chan))

    def forward(self, x):
        dim = self.dim
        right_ones = (dim + 1) if dim < 0 else (x.ndim - 1 - dim)
        gamma = self.gamma.reshape(-1, *((1,) * right_ones))
        return F.normalize(x, dim = dim) * (x.shape[dim] ** 0.5) * gamma

# feedforward

def shift_token(t):
    t, t_shift = t.chunk(2, dim = 1)
    t_shift = F.pad(t_shift, (0, 0, 0, 0, 1, -1), value = 0.)
    return torch.cat((t, t_shift), dim = 1)

class GEGLU(nn.Module):
    def forward(self, x):
        x, gate = x.chunk(2, dim = 1)
        return x * F.gelu(gate)

class FeedForward(nn.Module):
    def __init__(self, dim, mult = 4):
        super().__init__()

        inner_dim = int(dim * mult * 2 / 3)
        self.proj_in = nn.Sequential(
            nn.Conv3d(dim, inner_dim * 2, 1, bias = False),
            GEGLU()
        )

        self.proj_out = nn.Sequential(
            RMSNorm(inner_dim),
            nn.Conv3d(inner_dim, dim, 1, bias = False)
        )

    def forward(self, x, enable_time = True):

        is_video = x.ndim == 5
        enable_time &= is_video

        if not is_video:
            x = rearrange(x, 'b c h w -> b c 1 h w')

        x = self.proj_in(x)

        if enable_time:
            x = shift_token(x)

        out = self.proj_out(x)

        if not is_video:
            out = rearrange(out, 'b c 1 h w -> b c h w')

        return out

# best relative positional encoding

class ContinuousPositionBias(nn.Module):
    """ from https://arxiv.org/abs/2111.09883 """

    def __init__(
        self,
        *,
        dim,
        heads,
        num_dims = 1,
        layers = 2
    ):
        super().__init__()
        self.num_dims = num_dims

        self.net = nn.ModuleList([])
        self.net.append(nn.Sequential(nn.Linear(self.num_dims, dim), nn.SiLU()))

        for _ in range(layers - 1):
            self.net.append(nn.Sequential(nn.Linear(dim, dim), nn.SiLU()))

        self.net.append(nn.Linear(dim, heads))

    @property
    def device(self):
        return next(self.parameters()).device

    def forward(self, *dimensions):
        device = self.device

        shape = torch.tensor(dimensions, device = device)
        rel_pos_shape = 2 * shape - 1

        # calculate strides

        strides = torch.flip(rel_pos_shape, (0,)).cumprod(dim = -1)
        strides = torch.flip(F.pad(strides, (1, -1), value = 1), (0,))

        # get all positions and calculate all the relative distances

        positions = [torch.arange(d, device = device) for d in dimensions]
        grid = torch.stack(torch.meshgrid(*positions, indexing = 'ij'), dim = -1)
        grid = rearrange(grid, '... c -> (...) c')
        rel_dist = rearrange(grid, 'i c -> i 1 c') - rearrange(grid, 'j c -> 1 j c')

        # get all relative positions across all dimensions

        rel_positions = [torch.arange(-d + 1, d, device = device) for d in dimensions]
        rel_pos_grid = torch.stack(torch.meshgrid(*rel_positions, indexing = 'ij'), dim = -1)
        rel_pos_grid = rearrange(rel_pos_grid, '... c -> (...) c')

        # mlp input

        bias = rel_pos_grid.float()

        for layer in self.net:
            bias = layer(bias)

        # convert relative distances to indices of the bias

        rel_dist += (shape - 1)  # make sure all positive
        rel_dist *= strides
        rel_dist_indices = rel_dist.sum(dim = -1)

        # now select the bias for each unique relative position combination

        bias = bias[rel_dist_indices]
        return rearrange(bias, 'i j h -> h i j')

# helper classes

class Attention(nn.Module):
    def __init__(
        self,
        dim,
        dim_head = 64,
        heads = 8,
        flash = False,
        causal = False
    ):
        super().__init__()
        self.heads = heads
        self.scale = dim_head ** -0.5
        inner_dim = dim_head * heads

        self.attend = Attend(flash = flash, causal = causal)

        self.norm = RMSNorm(dim, dim = -1)

        self.to_q = nn.Linear(dim, inner_dim, bias = False)
        self.to_kv = nn.Linear(dim, inner_dim * 2, bias = False)
        self.to_out = nn.Linear(inner_dim, dim, bias = False)

        nn.init.zeros_(self.to_out.weight.data) # identity with skip connection

    def forward(
        self,
        x,
        rel_pos_bias = None
    ):
        x = self.norm(x)

        q, k, v = self.to_q(x), *self.to_kv(x).chunk(2, dim = -1)

        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), (q, k, v))

        out = self.attend(q, k, v, bias = rel_pos_bias)

        out = rearrange(out, 'b h n d -> b n (h d)')
        return self.to_out(out)

# main contribution - pseudo 3d conv

class PseudoConv3d(nn.Module):
    def __init__(
        self,
        dim,
        dim_out = None,
        kernel_size = 3,
        *,
        temporal_kernel_size = None,
        **kwargs
    ):
        super().__init__()
        dim_out = default(dim_out, dim)
        temporal_kernel_size = default(temporal_kernel_size, kernel_size)

        self.spatial_conv = nn.Conv2d(dim, dim_out, kernel_size = kernel_size, padding = kernel_size // 2)
        self.temporal_conv = nn.Conv1d(dim_out, dim_out, kernel_size = temporal_kernel_size, padding = temporal_kernel_size // 2) if kernel_size > 1 else None

        if exists(self.temporal_conv):
            nn.init.dirac_(self.temporal_conv.weight.data) # initialized to be identity
            nn.init.zeros_(self.temporal_conv.bias.data)

    def forward(
        self,
        x,
        enable_time = True
    ):
        b, c, *_, h, w = x.shape

        is_video = x.ndim == 5
        enable_time &= is_video

        if is_video:
            x = rearrange(x, 'b c f h w -> (b f) c h w')

        x = self.spatial_conv(x)

        if is_video:
            x = rearrange(x, '(b f) c h w -> b c f h w', b = b)

        if not enable_time or not exists(self.temporal_conv):
            return x

        x = rearrange(x, 'b c f h w -> (b h w) c f')

        x = self.temporal_conv(x)

        x = rearrange(x, '(b h w) c f -> b c f h w', h = h, w = w)

        return x

# factorized spatial temporal attention from Ho et al.

class SpatioTemporalAttention(nn.Module):
    def __init__(
        self,
        dim,
        *,
        dim_head = 64,
        heads = 8,
        add_feed_forward = True,
        ff_mult = 4,
        pos_bias = True,
        flash = False,
        causal_time_attn = False
    ):
        super().__init__()
        assert not (flash and pos_bias), 'learned positional attention bias is not compatible with flash attention'

        self.spatial_attn = Attention(dim = dim, dim_head = dim_head, heads = heads, flash = flash)
        self.spatial_rel_pos_bias = ContinuousPositionBias(dim = dim // 2, heads = heads, num_dims = 2) if pos_bias else None

        self.temporal_attn = Attention(dim = dim, dim_head = dim_head, heads = heads, flash = flash, causal = causal_time_attn)
        self.temporal_rel_pos_bias = ContinuousPositionBias(dim = dim // 2, heads = heads, num_dims = 1) if pos_bias else None

        self.has_feed_forward = add_feed_forward
        if not add_feed_forward:
            return

        self.ff = FeedForward(dim = dim, mult = ff_mult)

    def forward(
        self,
        x,
        enable_time = True
    ):
        b, c, *_, h, w = x.shape
        is_video = x.ndim == 5
        enable_time &= is_video

        if is_video:
            x = rearrange(x, 'b c f h w -> (b f) (h w) c')
        else:
            x = rearrange(x, 'b c h w -> b (h w) c')

        space_rel_pos_bias = self.spatial_rel_pos_bias(h, w) if exists(self.spatial_rel_pos_bias) else None

        x = self.spatial_attn(x, rel_pos_bias = space_rel_pos_bias) + x

        if is_video:
            x = rearrange(x, '(b f) (h w) c -> b c f h w', b = b, h = h, w = w)
        else:
            x = rearrange(x, 'b (h w) c -> b c h w', h = h, w = w)

        if enable_time:

            x = rearrange(x, 'b c f h w -> (b h w) f c')

            time_rel_pos_bias = self.temporal_rel_pos_bias(x.shape[1]) if exists(self.temporal_rel_pos_bias) else None

            x = self.temporal_attn(x, rel_pos_bias = time_rel_pos_bias) + x

            x = rearrange(x, '(b h w) f c -> b c f h w', w = w, h = h)

        if self.has_feed_forward:
            x = self.ff(x, enable_time = enable_time) + x

        return x

# resnet block

class Block(nn.Module):
    def __init__(
        self,
        dim,
        dim_out,
        kernel_size = 3,
        temporal_kernel_size = None,
        groups = 8
    ):
        super().__init__()
        self.project = PseudoConv3d(dim, dim_out, 3)
        self.norm = nn.GroupNorm(groups, dim_out)
        self.act = nn.SiLU()

    def forward(
        self,
        x,
        scale_shift = None,
        enable_time = False
    ):
        x = self.project(x, enable_time = enable_time)
        x = self.norm(x)

        if exists(scale_shift):
            scale, shift = scale_shift
            x = x * (scale + 1) + shift

        return self.act(x)

class ResnetBlock(nn.Module):
    def __init__(
        self,
        dim,
        dim_out,
        *,
        timestep_cond_dim = None,
        groups = 8
    ):
        super().__init__()

        self.timestep_mlp = None

        if exists(timestep_cond_dim):
            self.timestep_mlp = nn.Sequential(
                nn.SiLU(),
                nn.Linear(timestep_cond_dim, dim_out * 2)
            )

        self.block1 = Block(dim, dim_out, groups = groups)
        self.block2 = Block(dim_out, dim_out, groups = groups)
        self.res_conv = PseudoConv3d(dim, dim_out, 1) if dim != dim_out else nn.Identity()

    def forward(
        self,
        x,
        timestep_emb = None,
        enable_time = True
    ):
        assert not (exists(timestep_emb) ^ exists(self.timestep_mlp))

        scale_shift = None

        if exists(self.timestep_mlp) and exists(timestep_emb):
            time_emb = self.timestep_mlp(timestep_emb)
            to_einsum_eq = 'b c 1 1 1' if x.ndim == 5 else 'b c 1 1'
            time_emb = rearrange(time_emb, f'b c -> {to_einsum_eq}')
            scale_shift = time_emb.chunk(2, dim = 1)

        h = self.block1(x, scale_shift = scale_shift, enable_time = enable_time)

        h = self.block2(h, enable_time = enable_time)

        return h + self.res_conv(x)

# pixelshuffle upsamples and downsamples
# where time dimension can be configured

class Downsample(nn.Module):
    def __init__(
        self,
        dim,
        downsample_space = True,
        downsample_time = False,
        nonlin = False
    ):
        super().__init__()
        assert downsample_space or downsample_time

        self.down_space = nn.Sequential(
            Rearrange('b c (h p1) (w p2) -> b (c p1 p2) h w', p1 = 2, p2 = 2),
            nn.Conv2d(dim * 4, dim, 1, bias = False),
            nn.SiLU() if nonlin else nn.Identity()
        ) if downsample_space else None

        self.down_time = nn.Sequential(
            Rearrange('b c (f p) h w -> b (c p) f h w', p = 2),
            nn.Conv3d(dim * 2, dim, 1, bias = False),
            nn.SiLU() if nonlin else nn.Identity()
        ) if downsample_time else None

    def forward(
        self,
        x,
        enable_time = True
    ):
        is_video = x.ndim == 5

        if is_video:
            x = rearrange(x, 'b c f h w -> b f c h w')
            x, ps = pack([x], '* c h w')

        if exists(self.down_space):
            x = self.down_space(x)

        if is_video:
            x, = unpack(x, ps, '* c h w')
            x = rearrange(x, 'b f c h w -> b c f h w')

        if not is_video or not exists(self.down_time) or not enable_time:
            return x

        x = self.down_time(x)

        return x

class Upsample(nn.Module):
    def __init__(
        self,
        dim,
        upsample_space = True,
        upsample_time = False,
        nonlin = False
    ):
        super().__init__()
        assert upsample_space or upsample_time

        self.up_space = nn.Sequential(
            nn.Conv2d(dim, dim * 4, 1),
            nn.SiLU() if nonlin else nn.Identity(),
            Rearrange('b (c p1 p2) h w -> b c (h p1) (w p2)', p1 = 2, p2 = 2)
        ) if upsample_space else None

        self.up_time = nn.Sequential(
            nn.Conv3d(dim, dim * 2, 1),
            nn.SiLU() if nonlin else nn.Identity(),
            Rearrange('b (c p) f h w -> b c (f p) h w', p = 2)
        ) if upsample_time else None

        self.init_()

    def init_(self):
        if exists(self.up_space):
            self.init_conv_(self.up_space[0], 4)

        if exists(self.up_time):
            self.init_conv_(self.up_time[0], 2)

    def init_conv_(self, conv, factor):
        o, *remain_dims = conv.weight.shape
        conv_weight = torch.empty(o // factor, *remain_dims)
        nn.init.kaiming_uniform_(conv_weight)
        conv_weight = repeat(conv_weight, 'o ... -> (o r) ...', r = factor)

        conv.weight.data.copy_(conv_weight)
        nn.init.zeros_(conv.bias.data)

    def forward(
        self,
        x,
        enable_time = True
    ):
        is_video = x.ndim == 5

        if is_video:
            x = rearrange(x, 'b c f h w -> b f c h w')
            x, ps = pack([x], '* c h w')

        if exists(self.up_space):
            x = self.up_space(x)

        if is_video:
            x, = unpack(x, ps, '* c h w')
            x = rearrange(x, 'b f c h w -> b c f h w')

        if not is_video or not exists(self.up_time) or not enable_time:
            return x

        x = self.up_time(x)

        return x

# space time factorized 3d unet

class SpaceTimeUnet(nn.Module):
    def __init__(
        self,
        *,
        dim,
        channels = 3,
        dim_mult = (1, 2, 4, 8),
        self_attns = (False, False, False, True),
        temporal_compression = (False, True, True, True),
        resnet_block_depths = (2, 2, 2, 2),
        attn_dim_head = 64,
        attn_heads = 8,
        condition_on_timestep = True,
        attn_pos_bias = True,
        flash_attn = False,
        causal_time_attn = False
    ):
        super().__init__()
        assert len(dim_mult) == len(self_attns) == len(temporal_compression) == len(resnet_block_depths)
        num_layers = len(dim_mult)

        dims = [dim, *map(lambda mult: mult * dim, dim_mult)]
        dim_in_out = zip(dims[:-1], dims[1:])

        # determine the valid multiples of the image size and frames of the video

        self.frame_multiple = 2 ** sum(tuple(map(int, temporal_compression)))
        self.image_size_multiple = 2 ** num_layers

        # timestep conditioning for DDPM, not to be confused with the time dimension of the video

        self.to_timestep_cond = None
        timestep_cond_dim = (dim * 4) if condition_on_timestep else None

        if condition_on_timestep:
            self.to_timestep_cond = nn.Sequential(
                SinusoidalPosEmb(dim),
                nn.Linear(dim, timestep_cond_dim),
                nn.SiLU()
            )

        # layers

        self.downs = mlist([])
        self.ups = mlist([])

        attn_kwargs = dict(
            dim_head = attn_dim_head,
            heads = attn_heads,
            pos_bias = attn_pos_bias,
            flash = flash_attn,
            causal_time_attn = causal_time_attn
        )

        mid_dim = dims[-1]

        self.mid_block1 = ResnetBlock(mid_dim, mid_dim, timestep_cond_dim = timestep_cond_dim)
        self.mid_attn = SpatioTemporalAttention(dim = mid_dim, **attn_kwargs)
        self.mid_block2 = ResnetBlock(mid_dim, mid_dim, timestep_cond_dim = timestep_cond_dim)

        for _, self_attend, (dim_in, dim_out), compress_time, resnet_block_depth in zip(range(num_layers), self_attns, dim_in_out, temporal_compression, resnet_block_depths):
            assert resnet_block_depth >= 1

            self.downs.append(mlist([
                ResnetBlock(dim_in, dim_out, timestep_cond_dim = timestep_cond_dim),
                mlist([ResnetBlock(dim_out, dim_out) for _ in range(resnet_block_depth)]),
                SpatioTemporalAttention(dim = dim_out, **attn_kwargs) if self_attend else None,
                Downsample(dim_out, downsample_time = compress_time)
            ]))

            self.ups.append(mlist([
                ResnetBlock(dim_out * 2, dim_in, timestep_cond_dim = timestep_cond_dim),
                mlist([ResnetBlock(dim_in + (dim_out if ind == 0 else 0), dim_in) for ind in range(resnet_block_depth)]),
                SpatioTemporalAttention(dim = dim_in, **attn_kwargs) if self_attend else None,
                Upsample(dim_out, upsample_time = compress_time)

            ]))

        self.skip_scale = 2 ** -0.5 # paper shows faster convergence

        self.conv_in = PseudoConv3d(dim = channels, dim_out = dim, kernel_size = 7, temporal_kernel_size = 3)
        self.conv_out = PseudoConv3d(dim = dim, dim_out = channels, kernel_size = 3, temporal_kernel_size = 3)

    def forward(
        self,
        x,
        timestep = None,
        enable_time = True
    ):

        # some asserts

        assert not (exists(self.to_timestep_cond) ^ exists(timestep))
        is_video = x.ndim == 5

        if enable_time and is_video:
            frames = x.shape[2]
            assert divisible_by(frames, self.frame_multiple), f'number of frames on the video ({frames}) must be divisible by the frame multiple ({self.frame_multiple})'

        height, width = x.shape[-2:]
        assert divisible_by(height, self.image_size_multiple) and divisible_by(width, self.image_size_multiple), f'height and width of the image or video must be a multiple of {self.image_size_multiple}'

        # main logic

        t = self.to_timestep_cond(rearrange(timestep, '... -> (...)')) if exists(timestep) else None

        x = self.conv_in(x, enable_time = enable_time)

        hiddens = []

        for init_block, blocks, maybe_attention, downsample in self.downs:
            x = init_block(x, t, enable_time = enable_time)

            hiddens.append(x.clone())

            for block in blocks:
                x = block(x, enable_time = enable_time)

            if exists(maybe_attention):
                x = maybe_attention(x, enable_time = enable_time)

            hiddens.append(x.clone())

            x = downsample(x, enable_time = enable_time)

        x = self.mid_block1(x, t, enable_time = enable_time)
        x = self.mid_attn(x, enable_time = enable_time)
        x = self.mid_block2(x, t, enable_time = enable_time)

        for init_block, blocks, maybe_attention, upsample in reversed(self.ups):
            x = upsample(x, enable_time = enable_time)

            x = torch.cat((hiddens.pop() * self.skip_scale, x), dim = 1)

            x = init_block(x, t, enable_time = enable_time)

            x = torch.cat((hiddens.pop() * self.skip_scale, x), dim = 1)

            for block in blocks:
                x = block(x, enable_time = enable_time)

            if exists(maybe_attention):
                x = maybe_attention(x, enable_time = enable_time)

        x = self.conv_out(x, enable_time = enable_time)
        return x


import pandas as pd
import numpy as np
import os
import traceback

ticker = 'HSBC'
input_dir = r'../stock_data'
ticker_file = os.path.join(input_dir, ticker + '.csv')

def main():
    try:
        # read file
        df = pd.read_csv(ticker_file)
        df = pd.DataFrame(df)
        # simplify the table
        col1 = df['Weekday']
        col2 = df['Return']
        df = pd.DataFrame(df['Year'])
        df.insert(1, 'Weekday', col1)
        df.insert(2, 'Return', col2)

        print('opened file for ticker: ', ticker)

        # for each years calculate the result
        for i in range(2015, 2020):
            df_new = df[(df['Year']) == i]
            print("Year: ", i)

            # loop 3 times, 0: R; 1: R+; 2: R-
            for j in range(0, 3):
                df_temp = df_new
                if j == 1:
                    df_temp = df_new[(df_new['Return']) >= 0]
                    print("positive")
                elif j == 2:
                    df_temp = df_new[(df_new['Return']) < 0]
                    print("negative")
                # calculate the mean and std for each set, and count the number for each set
                calculator(df_temp)

                print()

    except Exception as e:
        print(e)
        print(traceback.format_exc())

def calculator(df):
    result = df.groupby(['Weekday'])['Return'].agg(['mean', 'std', 'count'])
    mean = df['Return'].mean()
    print(result, "\nAll year Mean ", mean)

main()


import sqlite3
from sqlite3 import Error
class DbConnection:
    def __init__(self, db_file):
        self.db_file = db_file
        self.connection = None

    def __enter__(self, *args, **kwargs):
        self.connection = create_connection(self.db_file)
        return self.connection

    def __exit__(self, *args, **kwargs):
        close_connection(self.connection)
        return

def create_connection(db_file):
    """ create a database connection to a SQLite database """
    connection = None
    try:
        connection = sqlite3.connect(db_file)
        return connection
    except Error as e:
        print(e)

    return connection

def close_connection(connection):
    if connection:
        connection.close()

def create_table(connection, create_table_sql):
    """ create a table from the create_table_sql statement
    :param connection: Connection object
    :param create_table_sql: a CREATE TABLE statement
    :return:
    """
    try:
        cursor = connection.cursor()
        cursor.execute(create_table_sql)
    except Error as e:
        print(e)

def create_connection_test(db_file):
    """ test connection to a SQLite database """
    connection = None
    try:
        connection = sqlite3.connect(db_file)
        print(sqlite3.version)
    except Error as e:
        print(e)
    finally:
        if connection:
            connection.close()


from flask_wtf import FlaskForm
from flask_wtf.file import FileField
from wtforms import StringField, SubmitField, BooleanField, IntegerField, FloatField
from wtforms.validators import DataRequired

class Equation(FlaskForm):
    equation = StringField('')
    submit = SubmitField()
class LineVariables(FlaskForm):
    slope = FloatField('Slope')
    intercept = FloatField('intercept')
    submit = SubmitField()

class BellVariables(FlaskForm):
    sq_a = FloatField('A')
    sq_b = FloatField('B')
    sq_c = FloatField('C')
    submit = SubmitField()

class DataFile(FlaskForm):
    file = FileField('File')
    submit = SubmitField('Submit')


import sys
from .rbt_node import RbtNode
from .structure import Structure
from .consts import BLACK, RED, NIL

class RbtStructure(Structure):
    def __init__(self):
        self.NIL_LEAF = RbtNode(0, NIL)
        self.root = self.NIL_LEAF
        super().__init__(self.root)

    def insert(self, key):
        self.count += 1
        key = self.prepare_word(key)
        node = RbtNode(key)
        node.parent = None
        node.val = key
        node.left = self.NIL_LEAF
        node.right = self.NIL_LEAF
        node.color = RED

        y = None
        x = self.root

        while x != self.NIL_LEAF:
            y = x
            if node.val < x.val:
                x = x.left
            else:
                x = x.right

        node.parent = y
        if y == None:
            self.root = node
        elif node.val < y.val:
            y.left = node
        else:
            y.right = node

        if node.parent == None:
            node.color = BLACK
            return

        if node.parent.parent == None:
            return

        self.__fix_insert(node)

    def min(self):
        if self.root:
            return self.min_node(self.root)
        else:
            return None

    def min_node(self, node):
        current = node
        while current != self.NIL_LEAF:
            if current.left == self.NIL_LEAF:
                break
            current = current.left

        return current

    def max(self):
        if self.root:
            return self.max_node(self.root)
        else:
            return None

    def max_node(self, node):
        current = node
        while current != self.NIL_LEAF:
            if current.right == self.NIL_LEAF:
                break
            current = current.right

        return current

    def __in_order_helper(self, node):
        if node != self.NIL_LEAF:
            self.__in_order_helper(node.left)
            sys.stdout.write(node.val + " ")
            self.__in_order_helper(node.right)

    def __search_tree_helper(self, node, key):
        if node == self.NIL_LEAF or key == node.val:
            return node

        self.find_counter += 1
        if key < node.val:
            return self.__search_tree_helper(node.left, key)
        return self.__search_tree_helper(node.right, key)

    def __iter__(self):
        if not self.root:
            return list()
        yield from self.root.__iter__()

    def __fix_delete(self, x):
        while x != self.root and x.color == BLACK:
            if x == x.parent.left:
                s = x.parent.right
                if s.color == RED:
                    s.color = BLACK
                    x.parent.color = RED
                    self.left_rotate(x.parent)
                    s = x.parent.right

                if s.left.color == BLACK and s.right.color == BLACK:
                    s.color = RED
                    x = x.parent
                else:
                    if s.right.color == BLACK:
                        s.left.color = BLACK
                        s.color = RED
                        self.right_rotate(s)
                        s = x.parent.right

                    s.color = x.parent.color
                    x.parent.color = BLACK
                    s.right.color = BLACK
                    self.left_rotate(x.parent)
                    x = self.root
            else:
                s = x.parent.left
                if s.color == RED:
                    s.color = BLACK
                    x.parent.color = RED
                    self.right_rotate(x.parent)
                    s = x.parent.left

                if s.left.color == BLACK and s.right.color == BLACK:
                    s.color = RED
                    x = x.parent
                else:
                    if s.left.color == BLACK:
                        s.right.color = BLACK
                        s.color = RED
                        self.left_rotate(s)
                        s = x.parent.left

                    s.color = x.parent.color
                    x.parent.color = BLACK
                    s.left.color = BLACK
                    self.right_rotate(x.parent)
                    x = self.root
        x.color = BLACK

    def __rb_transplant(self, u, v):
        if u.parent == None:
            self.root = v
        elif u == u.parent.left:
            u.parent.left = v
        else:
            u.parent.right = v
        v.parent = u.parent

    def delete(self, val):
        return self.__delete_node_helper(self.root, val)

    def __delete_node_helper(self, node, key):
        z = self.NIL_LEAF
        while node != self.NIL_LEAF:
            if node.val == key:
                z = node

            if key < node.val:
                node = node.left
            else:
                node = node.right

        if z == self.NIL_LEAF:
            return 0

        self.count -= 1
        y = z
        y_original_color = y.color
        if z.left == self.NIL_LEAF:
            x = z.right
            self.__rb_transplant(z, z.right)
        elif (z.right == self.NIL_LEAF):
            x = z.left
            self.__rb_transplant(z, z.left)
        else:
            y = self.min_node(z.right)
            y_original_color = y.color
            x = y.right
            if y.parent == z:
                x.parent = y
            else:
                self.__rb_transplant(y, y.right)
                y.right = z.right
                y.right.parent = y

            self.__rb_transplant(z, y)
            y.left = z.left
            y.left.parent = y
            y.color = z.color
        if y_original_color == 0:
            self.__fix_delete(x)

        return 1

    def  __fix_insert(self, k):
        while k.parent.color == RED:
            if k.parent == k.parent.parent.right:
                u = k.parent.parent.left
                if u.color == RED:
                    u.color = BLACK
                    k.parent.color = BLACK
                    k.parent.parent.color = RED
                    k = k.parent.parent
                else:
                    if k == k.parent.left:
                        k = k.parent
                        self.right_rotate(k)

                    k.parent.color = BLACK
                    k.parent.parent.color = RED
                    self.left_rotate(k.parent.parent)
            else:
                u = k.parent.parent.right

                if u.color == RED:
                    u.color = BLACK
                    k.parent.color = BLACK
                    k.parent.parent.color = RED
                    k = k.parent.parent
                else:
                    if k == k.parent.right:
                        k = k.parent
                        self.left_rotate(k)

                    k.parent.color = BLACK
                    k.parent.parent.color = RED
                    self.right_rotate(k.parent.parent)
            if k == self.root:
                break
        self.root.color = BLACK

    def inorder(self):
        self.__in_order_helper(self.root)

    def find(self, key):
        self.reset_find_counter()
        return self.__search_tree_helper(self.root, key)

    def successor(self, key):
        node = self.find(key)

        def successor_node(node):
            if node.right is not None and node.right.color != NIL:
                return self.min_node(node.right)

            succ = None
            root = self.root

            while root:
                if node.color == NIL:
                    break

                if node.val < root.val:
                    succ = root
                    root = root.left
                elif node.val > root.val:
                    root = root.right
                else:
                    break

            return succ
        if node:
            return successor_node(node)
        else:
            return None
    def left_rotate(self, x):
        y = x.right
        x.right = y.left
        if y.left != self.NIL_LEAF:
            y.left.parent = x

        y.parent = x.parent
        if x.parent == None:
            self.root = y
        elif x == x.parent.left:
            x.parent.left = y
        else:
            x.parent.right = y
        y.left = x
        x.parent = y

    def right_rotate(self, x):
        y = x.left
        x.left = y.right
        if y.right != self.NIL_LEAF:
            y.right.parent = x

        y.parent = x.parent
        if x.parent == None:
            self.root = y
        elif x == x.parent.right:
            x.parent.right = y
        else:
            x.parent.left = y
        y.right = x
        x.parent = y


"""
See `tts_calculation_without_bootstrapping` for a detailed explanation
on how to use the functions defined here.
"""

import numpy as np
import pandas as pd
from ps_from_rawdf import bv_prob, restrict_to_bvn
from tqdm import tqdm

def tts(prob: float, dt: float, desired: float = 0.99, unit: str = "log10") -> float:
    """Time to solution given the acquired probability, desired probability, and the time taken

    Args:
        prob (float): acquired probability
        dt (float): time for a single iteration
        desired (float, optional): desired probability. Defaults to 0.99.
        unit (str, optional): seconds or log10 seconds . Defaults to "log10".

    Returns:
        float: Time to solution
    """
    if prob == 0:
        return np.nan  # infinite TTS
    if unit == "default":
        return dt * (np.log(1 - desired) / np.log(1 - prob))
    if unit == "log10":
        return np.log10(tts(prob, dt, unit="default"))

def tts_on_bvn_row(row_data: pd.Series, seq: str, bv_durations: dict) -> float:
    """Given a row from a dataframe of success probabilities and the sequence
    we are concerned with, and the bv_durations compute the sucess probability
    for that particular entry.

    Args:
        row_data (pd.Series): a row from a dataframe of success probabilities
        seq (str): name of the DD sequence
        bv_durations (dict): dictionary representing durations for various BV-n

    Returns:
        float: Time to solution
    """
    bitstring = row_data.name  # get the bitstring like 00000
    no_of_1s = bitstring.count("1")
    duration = bv_durations[no_of_1s]
    prob = row_data[seq]
    # dur = row_data['circuit duration']
    return tts(prob, duration)

def tts_df(bv_n: pd.DataFrame, bv_durations: dict) -> pd.DataFrame:
    """Time to solution given success probabilities and corresponding durations

    Args:
        bv_n (pd.DataFrame): Success probabilities for all the oracles
        bv_durations (dict): Durations for all the oracles

    Returns:
        pd.DataFrame: Same dataframe as `bv_n` but probabilities are replaced with TTS
    """
    seqs = list(bv_n.columns)
    bv_n_tts = pd.DataFrame()
    for seq in seqs:
        bv_n_tts[seq] = bv_n.apply(
            lambda row: tts_on_bvn_row(row, seq, bv_durations), axis=1
        )
    return bv_n_tts

def tts_average(bv_n: pd.DataFrame, bv_durations: dict) -> pd.Series:
    """Averate Time to solution given a dataframe of TTS for each oracle

    Args:
        bv_n (pd.DataFrame): Dataframe with sucess probabilities for each oracle
        bv_durations (dict):  Durations for all the oracles

    Returns:
        pd.Series: Average TTS for each DD sequence
    """
    bv_n_tts = tts_df(bv_n, bv_durations)
    return bv_n_tts.apply(bv_prob)

def tts_average_for_all_n(bv_ns: list, bv_durations: dict) -> pd.DataFrame:
    """Given all the success probabilities and durations (for each oracle in all BV-n)
    return the average TTS for each BV-n

    Args:
        bv_ns (list): list of pd.DataFrame, each representing success probabilities for
            all oracles for a specific n
        bv_durations (dict): Durations for all the oracles

    Returns:
        pd.DataFrame: rows = n, columns = DD Sequence, entries are average TTS
    """

    tts_dataframe = pd.DataFrame()

    for n in range(0, len(bv_ns)):
        bv_n = bv_ns[n]
        tts_n = pd.DataFrame(
            tts_average(bv_n, bv_durations), columns=[n + 1]
        ).transpose()
        tts_dataframe = pd.concat([tts_dataframe, tts_n])

    return tts_dataframe

def tts_from_rawdf(rawdf: pd.DataFrame, bv_durations: dict) -> pd.DataFrame:
    """Time to solution for all sequences and BV-n given the rawdf and the duration.
    This is a composite function, that combines all the other functions in this script.

    Args:
        rawdf (pd.DataFrame): dataframes with prefix rawdf
        bv_durations (dict): Durations for all the oracles

    Returns:
        pd.DataFrame: rows = n, columns = DD Sequence, entries are average TTS
    """
    seqs = list(rawdf.keys()[0:3])  # first 3 elements are the sequences
    N = len(rawdf)  # 27
    ps_df_lst = []
    for n in tqdm(range(1, N)):
        ps_df_lst.append(restrict_to_bvn(rawdf, seqs, n=n))
    tts_all = tts_average_for_all_n(ps_df_lst, bv_durations)
    return tts_all


#reverse string

def reverse_string(string):
    reversed_list = list(reversed(string))
    reversed_string = ''.join(reversed_list)
    return reversed_string

print(reverse_string('ab123'))

# reverse list

lst = [1, 2, 3]
lst.reverse()
print(lst)

# REVERSE RETURNS NONE !!!

lst = [1, 2, 3]
rev_list = lst.reverse()
print(rev_list) # WILL PRINT NONE


import torch
import string
import os
import torch.nn as nn

class TransformerModel(nn.Module):
    def __init__(self, chars, embed_size, heads, num_layers, hidden_size, sequence_length, lr, epochs, checkpoint_interval, clip_value):
        super(TransformerModel, self).__init__()
        self.chars = chars
        self.vocab_size = len(chars)
        self.embed_size = embed_size
        self.heads = heads
        self.num_layers = num_layers
        self.hidden_size = hidden_size
        self.sequence_length = sequence_length
        self.lr = lr
        self.epochs = epochs
        self.checkpoint_interval = checkpoint_interval
        self.clip_value = clip_value

        self.embedding = nn.Embedding(self.vocab_size, embed_size)
        self.transformer = nn.Transformer(embed_size, nhead=heads, num_encoder_layers=num_layers, num_decoder_layers=num_layers)
        self.fc = nn.Linear(embed_size, self.vocab_size)

    def forward(self, x):
        x = self.embedding(x)
        x = self.transformer(x, x)
        x = self.fc(x)
        return x

def char_to_index(char, model):
    return model.chars.index(char)

def index_to_char(index, model):
    return model.chars[index]

def sample_with_temperature(logits, temperature):
    probs = torch.nn.functional.softmax(logits / temperature, dim=0)
    return torch.multinomial(probs, 1).item()

# Load checkpoint
checkpoint_path = "model_checkpoint.pth"
if not os.path.exists(checkpoint_path):
    print("Checkpoint file not found. Please train the model first.")
    exit()

checkpoint = torch.load(checkpoint_path)
hyperparameters = checkpoint['hyperparameters']

# First, we create a dummy list of chars to instantiate the model
dummy_chars = ['a']  # This is just a placeholder

model = TransformerModel(
    chars=checkpoint['chars'],
    embed_size=hyperparameters['embed_size'],
    heads=hyperparameters['heads'],
    num_layers=hyperparameters['num_layers'],
    hidden_size=hyperparameters['hidden_size'],
    sequence_length=hyperparameters['sequence_length'],
    lr=hyperparameters['lr'],
    epochs=hyperparameters['epochs'],
    checkpoint_interval=hyperparameters['checkpoint_interval'],
    clip_value=hyperparameters['clip_value']
)

# Now load the state dict.
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()

print("Chat with the model! Type 'exit' to end the chat.")
temperature = 0.5
output_length = 200

while True:
    input_text = input("> ").lower()
    if input_text == 'exit':
        break

    char_indices = [char_to_index(char, model) for char in input_text]  # Pass model to char_to_index
    input_tensor = torch.tensor([char_indices])

    # Generate additional characters after the user input
    for _ in range(output_length):
        output = model(input_tensor)
        char_idx = sample_with_temperature(output[0, -1], temperature)
        output_char = index_to_char(char_idx, model)  # Pass model to index_to_char
        print(output_char, end='', flush=True)
        char_indices.append(char_idx)
        input_tensor = torch.tensor([char_indices])

    print()  # to move to a new line after the generated sequence

print("Chat ended.")


import os
from os.path import exists
from glob import glob
import path_define

CRAWLEDLISTPATH= path_define.get_CRAWLEDLISTPATH()+"\crawled_list.txt"
ERROR_FILE_PATH= path_define.get_ERROR_FILE_PATH()
CRAWLED_HTML_PATH = path_define.get_CRAWLED_HTML_PATH()

def remove_files():
    files=glob(CRAWLED_HTML_PATH+'\*.html')
    for filename in files:
        os.remove(filename)
    if exists(CRAWLEDLISTPATH):
        os.remove(CRAWLEDLISTPATH)
    if exists(ERROR_FILE_PATH):
        os.remove(ERROR_FILE_PATH)


import pandas
import numpy as np
from sklearn.cluster import KMeans
from gestion.gestion_des_fichiers.afficheur import Afficheur
from gestion.elements_fichiers.data_base import Data_Base
from gestion.elements_fichiers.pays import Pays

class Resume(Afficheur):
	"""La classe Resume permet de renvoyer des data frames"""

	def table_criteres(self, liste_pays_a_afficher):
		"""Cette méthode permet de créer un data frame sur les pays choisis au préalable
		par la méthode critère_usuels à l'aide de la bibliothèque pandas.
		Valeurs_pays est un tableau qui prend toutes les informations chiffrées des différents pays. """

		noms_pays = [self.simplification(pays.get_name()) for pays in liste_pays_a_afficher]

		criteres = ['Superficie', 'Population', 'Croissance démographique', 'Inflation', 'Dette', 'Taux de chômage', 'Taux de dépenses en santé', 'Taux de dépenses en éducation', 'Taux de dépenses militaires']

		valeurs_pays = [[self.simplification(pays.get_superficie()) for pays in liste_pays_a_afficher],
						[self.simplification(pays.get_pop()) for pays in liste_pays_a_afficher],
						[self.simplification(pays.get_croissance_demo()) for pays in liste_pays_a_afficher],
						[self.simplification(pays.get_inflation()) for pays in liste_pays_a_afficher],
						[self.simplification(pays.get_dette()) for pays in liste_pays_a_afficher],
						[self.simplification(pays.get_chomage()) for pays in liste_pays_a_afficher],
						[self.simplification(pays.get_depenses_sante()) for pays in liste_pays_a_afficher],
						[self.simplification(pays.get_depenses_education()) for pays in liste_pays_a_afficher],
						[self.simplification(pays.get_depenses_militaires()) for pays in liste_pays_a_afficher]]

		return pandas.DataFrame(valeurs_pays, index = criteres, columns = noms_pays)

	def top_and_flop(self, critere, taille_classement):
		"""Cette méthode renvoie deux data frames.
		Le premier correspond aux 'taille_classement'(entier) premiers pays selon un certain critère.
		Le second data frame correspond aux 'taille_classement'(entier) dernier pays selon un certain critère."""

		donnees = Data_Base().donnees

		liste_triee_selon_critere = self.liste_triee_selon_critere(donnees, critere)

		top = []
		flop = []
		rang_top = range(1,taille_classement+1)
		rang_flop = [len(liste_triee_selon_critere)-i for i in range(taille_classement)]
		col = ['PAYS', critere.upper()]
		for i in range(taille_classement):
				top.append([liste_triee_selon_critere[-i-1][1].get_name(), self.simplification(liste_triee_selon_critere[-i-1][1].get_critere(critere))])
				flop.append([liste_triee_selon_critere[i][1].get_name(), self.simplification(liste_triee_selon_critere[i][1].get_critere(critere))])

		return (pandas.DataFrame(top, index = rang_top, columns = col), pandas.DataFrame(flop, index = rang_flop, columns = col))

	def sup_inf_seuil(self, critere, seuil):
		"""Cette méthode crée deux tableaux. Elle prend en argument un critère et un seuil.
		La premier tableau correspond aux pays dépassant le seuil déterminé par l'utilisateur.
		Le second correspond aux pays sous ce seuil. Elle renvoie 4 informations :
		le nombre de pays dépassant le seuil avec le tableau des pays correspondant avec les valeurs de ce critère,
		et le nombre de pays ne dépassant pas le seuil avec le tableau des pays correspondant avec les valeurs de ce critère"""

		donnees = Data_Base().donnees

		liste_triee_selon_critere = self.liste_triee_selon_critere(donnees, critere)
		liste_triee_selon_critere.reverse()
		pays_sup_seuil = []
		pays_inf_seuil = []
		col = ['PAYS', critere.upper()]

		for elm in liste_triee_selon_critere:
			if elm[0] >= seuil:
				pays_sup_seuil.append([elm[1].get_name(), self.simplification(elm[1].get_critere(critere))])
			else :
				pays_inf_seuil.append([elm[1].get_name(), self.simplification(elm[1].get_critere(critere))])

		tableau_pays_sup = pandas.DataFrame(pays_sup_seuil, index=[i+1 for i in range(len(pays_sup_seuil))], columns=col)
		tableau_pays_inf = pandas.DataFrame(pays_inf_seuil, index=[i+1 for i in range(len(pays_inf_seuil))], columns=col)

		return (len(pays_sup_seuil), tableau_pays_sup, len(pays_inf_seuil), tableau_pays_inf)

	def tableau_classes_age(self, liste_pays_a_afficher):
		"""Cette méthode renvoie un data frame concernant les différentes classes d'âges
		en prenant en compte les pays choisis au préalable."""

		noms_pays = [self.simplification(pays.get_name()) for pays in liste_pays_a_afficher]

		criteres = ["0-14", "15-24", "25-54", "55-64", ">= 65"]

		valeurs_classes_age = [[self.simplification(pays.get_classe_age1()) for pays in liste_pays_a_afficher],
							[self.simplification(pays.get_classe_age2()) for pays in liste_pays_a_afficher],
							[self.simplification(pays.get_classe_age3()) for pays in liste_pays_a_afficher],
							[self.simplification(pays.get_classe_age4()) for pays in liste_pays_a_afficher],
							[self.simplification(pays.get_classe_age5()) for pays in liste_pays_a_afficher]]

		return pandas.DataFrame(valeurs_classes_age, index = criteres, columns = noms_pays)

	def somme(self):
		"""Cette méthode renvoie un data frame d'une colonne. Elle fait la somme
		des différents critères sur les pays."""

		donnees = Data_Base().donnees
		superficie_tot = sum([elm[0] for elm in self.liste_triee_selon_critere(donnees, 'superficie')])
		population_tot = sum([elm[0] for elm in self.liste_triee_selon_critere(donnees, 'population')])
		dette_tot = sum([elm[0] for elm in self.liste_triee_selon_critere(donnees, 'dette')])

		criteres = ["La somme des superficies de tous les pays", "La somme des populations de tous les pays", "La somme des dettes de tous les pays"]
		index = criteres
		valeurs = [[superficie_tot],[population_tot],[dette_tot]]
		return pandas.DataFrame(valeurs, index = criteres, columns = [''])

	def summary(self):
		"""La méthode summary renvoie un data frame donnant des statistiques sur les différents critères (nombres de valeurs, moyenne, écart-type, ...)
		mais uniquement sur les colonnes numériques. Pour cela, la méthode crée un dictionnaire où chaque clé correspond à un critère
		que l'on associe à un data frame. Enfin, on utilise la méthode describe de pandas pour avoir notre résumé statistique."""

		donnees = Data_Base().donnees

		superficie = pandas.Series([elm[0] for elm in self.liste_triee_selon_critere(donnees, 'superficie')])
		population = pandas.Series([elm[0] for elm in self.liste_triee_selon_critere(donnees, 'population')])
		croissance_demographique = pandas.Series([elm[0] for elm in self.liste_triee_selon_critere(donnees, 'croissance démographique')])
		inflation = pandas.Series([elm[0] for elm in self.liste_triee_selon_critere(donnees, 'inflation')])
		dette = pandas.Series([elm[0] for elm in self.liste_triee_selon_critere(donnees, 'dette')])
		chomage = pandas.Series([elm[0] for elm in self.liste_triee_selon_critere(donnees, 'chômage')])
		depenses_sante = pandas.Series([elm[0] for elm in self.liste_triee_selon_critere(donnees, 'dépenses santé')])
		depenses_éducation = pandas.Series([elm[0] for elm in self.liste_triee_selon_critere(donnees, 'dépenses éducation')])
		depenses_militaires = pandas.Series([elm[0] for elm in self.liste_triee_selon_critere(donnees, 'dépenses militaires')])

		dic = {'Superficie' : superficie,
			'Population' : population,
			'Croissance démo' : croissance_demographique,
			'Inflation' : inflation,
			'Dette' : dette,
			'Chômage' : chomage,
			'Dépenses santé' : depenses_sante,
			'Dépenses éducation' : depenses_éducation,
			'Dépenses militaires' : depenses_militaires}

		return pandas.DataFrame(dic).describe()

	def clustering(self, nb_clusters):
		"""Cette méthode utilise la méthode des Kmeans pour permettre à la fonction de créer des classes entre les pays.
		Elle prend en argument le nombre de classes que l'utilisateur veut former.
		Cette méthode renvoie un data frame où chaque colonne correspond à chaque classe
		avec le nom des pays dans celles ci."""

		donnees = Data_Base().donnees

		criteres = ['superficie', 'population', 'croissance démographique', 'inflation', 'dette', 'chômage', 'dépenses santé', 'dépenses éducation', 'dépenses militaires']

		valeurs_pays = []
		liste_des_pays_comptabilises = []

		for num_pays in range(len(donnees)):
			pays = Pays(num_pays, donnees)
			if pays.get_name():
				if True not in [self.numerisation_critere(pays, critere) == 'NA' for critere in criteres]:
					liste_des_pays_comptabilises.append(self.simplification(pays.get_name()))
					valeurs_pays.append([self.numerisation_critere(pays, critere) for critere in criteres])

		clf = KMeans(n_clusters=nb_clusters)
		clf.fit(np.array(valeurs_pays))
		labels = clf.labels_

		pays_dans_clusters = [[] for i in range(nb_clusters)]

		for i in range(len(labels)):
			for cl in range(nb_clusters):
				if labels[i] == cl:
					pays_dans_clusters[cl].append(liste_des_pays_comptabilises[i])

		for cluster in pays_dans_clusters:
			while len(cluster) < max([len(cl) for cl in pays_dans_clusters]):
				cluster.append('')

		liste_des_clusters = ['Cluster {}'.format(i+1) for i in range(nb_clusters)]

		return pandas.DataFrame(pays_dans_clusters, index = liste_des_clusters).T


import os
import time
import random

print('input the mode for training mode')
mode = int(input())

'''
1. original == 0:
  name = 'Random'

2. original == 1:
  name = 'PatDNN'

3. original == 2:
  name = 'Pconv'

The model can be changed with 'ResNet-Q, WRN16-4'

entries: the number of non-zero weights.
If entries is 8, only 1 weight is zero.
'''

if mode == 1 : #
    original = [0,1]
    epoch = [100]
    wb = [1]
    ab = [4]
    seed_candidate = [1106, 1992]
    model = 'ResNet_Q'
    dataset = ['cifar10']
    GPU = 0
    entries = [4,5,6,7,8]

elif mode == 2 : #
    original = [2]
    epoch = [100]
    wb = [1]
    ab = [4]
    seed_candidate = [1106, 1992]
    model = 'ResNet_Q'
    dataset = ['cifar10']
    GPU = 0
    entries = [4]

elif mode == 5 : #
    original = [0]
    epoch = [200]
    wb = [2]
    ab = [8]
    seed_candidate = [1106, 1992]
    model = 'Wide_ResNet_Q'
    dataset = ['cifar100']
    GPU = 3
    entries = [4,5,6,7,8]
#########################################

for seed in seed_candidate :
    for wbb in wb :
        for abb in ab :
            for dset in dataset:
                for epo in epoch :
                    for ori in original :
                        for en in entries:
                            os.system( 'python3 utilization_pattern.py' + ' --dataset ' + str(dset) + ' --model ' +str(model) + ' --original '+ str(ori)
                            + ' --GPU ' + str(GPU)  + ' --epoch ' + str(epo) + ' --ab ' + str(abb) + ' --wb ' + str(wbb) + ' --seed ' + str(seed) + ' --entry ' + str(en))

time.sleep(10)


import socket
import HMAC
import Elgamal
from random import randint as rint

# FILE NOT USEFUL, MADE FOR TESTING

port = 5000

q = Elgamal.q
alpha = Elgamal.alpha
public_key_A = []
random_int = rint(1, q)

message_to_sent = '-__-, HEY, HELLO'

print("I AM SENDING THIS MESSAGE :", message_to_sent)

s = socket.socket()
s.bind(('', port))
s.listen()
c, addr = s.accept()
message = c.recv(1024)
message = message.decode()
public_key_A = list(map(int, message.split(" ")))

K = pow(public_key_A[2], random_int, q)
encrypted_message = Elgamal.encrypt(q, alpha, random_int, K, message_to_sent)
print(encrypted_message)
c.send(str.encode(encrypted_message))
s.close()


# -*- coding: UTF-8 -*-
from app import config
import os, sys, stat
import ipaddress
import hashlib

class Utils(object):

    @staticmethod
    def get_network_settings(ipaddr):
        addr4 = ipaddress.ip_address(ipaddr)
        gw = None
        nm = None
        found = False

        for subnet in config.local_networks:
            if addr4 in ipaddress.ip_network(subnet):
                gw = config.local_networks[subnet]['gateway']
                nm = config.local_networks[subnet]['netmask']
                pr = config.local_networks[subnet]['prefix']
                found = True

        if not found:
            raise Exception("NotFound", "Не найдены настройки сети для указанного IP-адреса {}".format(str(addr4)))

        return [gw, nm, pr]

    @staticmethod
    def get_dhcp_addr(ipaddr):
        addr4 = ipaddress.ip_address(ipaddr) + 1
        return str(addr4)

    @staticmethod
    def get_token():
        return hashlib.md5(os.urandom(32)).hexdigest()

    @staticmethod
    def create_preseed_conf(adman_id, params, token):
        preseed_template = '{preseeddir}/{os}{osver}_{diskpart}.seed'.format(
            preseeddir=config.utils['preseed-directory'],
            os=params['os'].lower(),
            osver=params['osver'].lower(),
            diskpart=params['diskpart'].lower()
        )

        preseed_config = '{confdir}/s{srv}/s{srv}.seed'.format(
            confdir=config.utils['config-directory'],
            srv=adman_id
        )

        mkconfig_cmd = '/bin/cp -f {src} {dst}'.format(
            src=preseed_template,
            dst=preseed_config
        )

        stream = os.popen(mkconfig_cmd)
        output = stream.read()

        # -- Fill template

        gateway, netmask, _ = Utils.get_network_settings(params['ipaddr'])

        # Read in the file
        with open(preseed_config, 'r') as file:
            filedata = file.read()

        # Replace the target string
        filedata = filedata.replace('__ipaddr__', str(params['ipaddr']))
        filedata = filedata.replace('__netmask__', netmask)
        filedata = filedata.replace('__gateway__', gateway)
        filedata = filedata.replace('__srv_name__', str(adman_id))
        filedata = filedata.replace('__passwdhash__', str(params['passwdhash']))
        filedata = filedata.replace('__token__', str(token))
        filedata = filedata.replace('__boot_host__', str(config.utils['boot_host']))

        # Write the file out again
        with open(preseed_config, 'w') as file:
            file.write(filedata)

        os.chmod(preseed_config, 436)

    @staticmethod
    def create_server_dir(server_id):
        create_cmd = "/bin/mkdir -p {dir}/s{server_id}".format(
            dir=config.utils['config-directory'],
            server_id=server_id
        )

        stream = os.popen(create_cmd)
        output = stream.read()
        return output

    @staticmethod
    def remove_symbol_link(link):
        rm_cmd = "/bin/rm -rf {dir}/{link}".format(
            dir = config.utils['config-directory'],
            link=link
        )

        stream = os.popen(rm_cmd)
        output = stream.read()
        return output

    @staticmethod
    def create_symbol_link(server_id, link):
        os.chdir(config.utils['config-directory'])
        create_cmd = "/bin/ln -s s{server_id} {link}".format(
            server_id=server_id,
            link=link
        )

        stream = os.popen(create_cmd)
        output = stream.read()
        return output

    @staticmethod
    def create_config(args, adman_id, token=''):

        diskpart = int(args['diskpart'])
        conf_dir = str(config.utils['config-directory'])
        url = ""

        if diskpart == 2:
            url = str(config.utils['iscsi_url'])
        else:
            url = str(config.utils['boot_url'])

        ipxe_config = '{confdir}/s{srv}/boot.ipxe'.format(
            confdir=conf_dir,
            srv=adman_id
        )

        tpl_key = "{os}_{osver}".format(os = args['os'].lower(), osver = args['osver'].lower())
        if diskpart == 2 and args['os'].lower() == 'local':
            tpl_key = "{os}_{osver}".format(os=args['os'].lower(), osver=2)

        tpl = config.templates[tpl_key]

        with open(ipxe_config, 'a') as file:
            file.truncate(0)
            file.writelines(tpl)

        # Read in the file
        with open(ipxe_config, 'r') as file:
            filedata = file.read()

        # Replace the target string
        filedata = filedata.replace('__srv_name__', str(adman_id))
        filedata = filedata.replace('__boot_url__', url)
        filedata = filedata.replace('__token__', str(token))

        # Write the file out again
        with open(ipxe_config, 'w') as file:
            file.write(filedata)

        os.chmod(ipxe_config, 436)

    @staticmethod
    def create_dhcp_config(adman_id, mac_addr, dhcp_addr, router=True):

        conf_dir = "/srv/hosts"
        dhcp_config = '{confdir}/s{srv}.conf'.format(
            confdir=conf_dir,
            srv=adman_id
        )

        with open(dhcp_config, 'a') as file:
            file.truncate(0)
            port_number = 0
            tpl_list = []
            for mac in mac_addr:
                tpl_list.append('#Autogenerated DHCP configuration\n')
                tpl_list.append('  host s{srv_name}p{port_number} {{\n'.format(srv_name=adman_id, port_number=port_number))
                tpl_list.append('    hardware ethernet {mac_addr};\n'.format(mac_addr=mac.mac_addr))
                tpl_list.append('    fixed-address {dhcp_addr};\n'.format(dhcp_addr=dhcp_addr))
                if router:
                    tpl_list.append('    option routers 10.0.224.4;\n')
                tpl_list.append('  }\n\n')
                port_number += 1

            file.writelines(tpl_list)

        os.chmod(dhcp_config, 436)

    @staticmethod
    def create_install_bat(args, adman_id):
        bat_config = '{confdir}/s{srv}/install.bat'.format(
            confdir=config.utils['config-directory'],
            srv=adman_id
        )

        osver = args['osver'].lower()

        with open(bat_config, 'a') as file:
            file.truncate(0)
            file.write('wpeinit\n')
            file.write('ping 10.0.222.1\n')
            file.write('ping 10.0.222.1\n')
            file.write('ping 10.0.222.1\n')
            file.write('net use j: \\\\__boot_host__\\images\\__os_ver__\\amd64 /user:user pass\n')
            file.write('net use k: \\\\__boot_host__\\configfile /user:user pass\n')
            file.write('j:\\setup.exe /unattend:k:\\s__srv_name__\\s__srv_name__.seed\n')

        # Read in the file
        with open(bat_config, 'r') as file:
            filedata = file.read()

        # Replace the target string
        filedata = filedata.replace('__srv_name__', str(adman_id))
        filedata = filedata.replace('__os_ver__', str(osver))
        filedata = filedata.replace('__boot_host__', str(config.utils['boot_host']))

        # Write the file out again
        with open(bat_config, 'w') as file:
            file.write(filedata)

        os.chmod(bat_config, 436)

    @staticmethod
    def create_win_set_ip_ps(ipaddr, token, adman_id):
        ps_config = '{confdir}/s{srv}/set-ip.ps1'.format(
            confdir=config.utils['config-directory'],
            srv=adman_id
        )

        with open(ps_config, 'a') as file:
            file.truncate(0)
            file.write('$ifaceindex = get-wmiobject win32_networkadapter -filter "netconnectionstatus = 2" | select -expand InterfaceIndex; ')
            file.write('netsh interface ipv4 set address name=$ifaceindex source=static addr=__ip_addr__  mask=__netmask__ gateway=__gateway__ gwmetric=1; ')
            file.write('netsh interface ipv4 add dns $ifaceindex 8.8.8.8 index=1; ')
            file.write('Timeout /T 15; ')
            file.write('ping 8.8.8.8; ')
            file.write('[System.Net.WebRequest]::Create("http://install.adman.cloud/api/v1.0/install/complete/__srv_name__?token=__token__").GetResponse(); ')

        gateway, netmask, _ = Utils.get_network_settings(ipaddr)
        # Read in the file
        with open(ps_config, 'r') as file:
            filedata = file.read()

        # Replace the target string
        filedata = filedata.replace('__ip_addr__', str(ipaddr))
        filedata = filedata.replace('__netmask__', netmask)
        filedata = filedata.replace('__gateway__', gateway)
        filedata = filedata.replace('__srv_name__', str(adman_id))
        filedata = filedata.replace('__token__', token)

        # Write the file out again
        with open(ps_config, 'w') as file:
            file.write(filedata)

        os.chmod(ps_config, 436)

    @staticmethod
    def remove_install_bat(adman_id):
        bat_config = '{confdir}/s{srv}/install.bat'.format(
            confdir=config.utils['config-directory'],
            srv=adman_id
        )

        if os.path.exists(bat_config):
            os.remove(bat_config)

    @staticmethod
    def remove_preseed_conf(adman_id):
        preseed_config = '{confdir}/s{srv}/s{srv}.seed'.format(
            confdir=config.utils['config-directory'],
            srv=adman_id
        )

        if os.path.exists(preseed_config):
            os.remove(preseed_config)

    @staticmethod
    def remove_win_set_ip_ps(adman_id):
        ps_config = '{confdir}/s{srv}/set-ip.ps1'.format(
            confdir=config.utils['config-directory'],
            srv=adman_id
        )

        if os.path.exists(ps_config):
            os.remove(ps_config)


import copy
import re
import time

DOTS = [
    ['ａ', 'ｂ', 'ｃ', 'ｄ', 'ｅ', ],
    ['ｆ', 'ｇ', 'ｈ', 'ｉ', 'ｊ', ],
    ['ｋ', 'ｌ', 'ｍ', 'ｎ', 'ｏ', ],
    ['ｐ', 'ｑ', 'ｒ', 'ｓ', 'ｔ', ],
    ['ｕ', 'ｖ', 'ｗ', 'ｘ', 'ｙ', ],
]

CELLS = [
    [0, 0, 0, 0, ],
    [0, 0, 0, 0, ],
    [0, 0, 0, 0, ],
    [0, 0, 0, 0, ],
]

ROWS = [
    [0, 0, 0, 0, ],
    [0, 0, 0, 0, ],
    [0, 0, 0, 0, ],
    [0, 0, 0, 0, ],
    [0, 0, 0, 0, ],
]

COLS = [
    [0, 0, 0, 0, ],
    [0, 0, 0, 0, ],
    [0, 0, 0, 0, ],
    [0, 0, 0, 0, ],
    [0, 0, 0, 0, ],
]

# ターミナルの文字色・背景色・太さなど
class Color:
    BLACK = '\033[30m'
    RED = '\033[31m'
    GREEN = '\033[32m'
    YELLOW = '\033[33m'
    BLUE = '\033[34m'
    PURPLE = '\033[35m'
    CYAN = '\033[36m'
    WHITE = '\033[37m'
    BG_BLACK = '\033[40m'
    BG_RED = '\033[41m'
    BG_GREEN = '\033[42m'
    BG_YELLOW = '\033[43m'
    BG_BLUE = '\033[44m'
    BG_PURPLE = '\033[45m'
    BG_CYAN = '\033[46m'
    BG_WHITE = '\033[47m'
    END = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'
    INVISIBLE = '\033[08m'
    REVERCE = '\033[07m'

# idは変更不可
class PLAYER1:
    id = 'player1'
    name = 'プレイヤー1'
    color = Color.YELLOW
    bg_color = Color.BG_YELLOW
    number = 1
    count = 0

# idは変更不可
class PLAYER2:
    id = 'player2'
    name = 'プレイヤー2'
    color = Color.PURPLE
    bg_color = Color.BG_PURPLE
    number = 2
    count = 0

ROW_LINE = 'ーー'
COL_LINE = '｜'
HR = '------------------------------------------'
PADDING_LEFT = '　　　　'

def main():
    # ゲーム開始アナウンス
    print(Color.BOLD + 'ドットアンドボックスへようこそ！\nゲームを開始します' + Color.END)
    # 2秒
    time.sleep(2)

    retry = True
    retry_text_color = Color.BOLD
    while retry:
        dotAndBox()
        retry = yes_no_input('もう一度遊ぶ？\n', retry_text_color)

def dotAndBox():
    # 初期化
    cells = copy.deepcopy(CELLS)
    rows = copy.deepcopy(ROWS)
    cols = copy.deepcopy(COLS)
    current_player = None
    player1 = PLAYER1
    player2 = PLAYER2
    current_player = togglePlayer(current_player, player1, player2)
    # 現ターンでマスが埋められたかどうか
    filled = False

    display(cells, rows, cols, [player1, player2])

    while True:
        print(current_player.color + current_player.name + 'のターン！' + Color.END)

        #　入力値から始点と終点を取得
        start_xy, end_xy = getInput()

        # データを更新する
        cells, rows, cols, start_xy, end_xy, player, filled = update(cells, rows, cols, start_xy, end_xy, current_player, filled)

        # 結果を表示する
        display(cells, rows, cols, [player1, player2])

        if isEnd(cells):
            # 終了
            print('ゲーム終了！')
            winner = player1
            if player1.count == player2.count:
                print('引き分け！')
                break

            if winner.count < player2.count:
                winner = player2
            print(winner.color + winner.name + 'の勝ち！' + Color.END + '\n')

            # リセット
            player1.count = 0
            player2.count = 0

            break
        else:
            # 継続
            if not filled:
                current_player = togglePlayer(current_player, player1, player2)

# プレイヤーのターンをチェンジ
def togglePlayer(current_player, player1, player2):
    if current_player is None:
        return player1
    elif current_player.id == 'player1':
        return player2
    elif current_player.id == 'player2':
        return player1
    else:
        raise Exception('現在ターンのプレイヤーが不適切です。')

# 入力値を取得する
def getInput():
    # 始点を入力
    while True:
        start_input = input("始点を入力してください\n>>")
        start_xy = getCoordinates(start_input)

        if not start_xy:
            # 再入力
            continue

        # 問題がなければ次の処理へ進む
        break

    # 終点を入力
    while True:
        end_input = input("終点を入力してください\n(zで始点の入力に戻る)\n>>")

        # zが入力されたら始点の入力に戻る
        if end_input == 'z':
            print(Color.RED + '再入力します' + Color.END)
            start_xy, end_xy = getInput()
            return start_xy, end_xy
        end_xy = getCoordinates(end_input)

        if not end_xy:
            # 再入力
            continue

        # 問題がなければ次の処理へ進む
        break

    # 始点と終点の関係を検証
    while not validateStartAndEndXY(start_xy, end_xy):
        # 再入力
        start_xy, end_xy = getInput()

    return start_xy, end_xy

# 座標を取得
def getCoordinates(point):
    # 入力値の妥当性を検証
    if not validateInput(point):
        return False

    # 半角英->全角英
    point = point.translate(str.maketrans(
        {chr(0x0021 + i): chr(0xFF01 + i) for i in range(94)}))

    # DOTS配列から座標を取得
    for y, dot_row in enumerate(DOTS):
        for x, dot in enumerate(dot_row):
            if dot == point:
                return [x, y]

    # DOTS配列に該当する値がなければ真偽値Falseを戻す
    return False

# 入力された値の妥当性を検証し、真偽値を返す

def validateInput(input):
    # 半角a-yであるかどうか
    if not re.match(r'^[a-y]$', input):
        print(Color.RED + 'aからyまでの半角英数を入力してください' + Color.END)
        return False
    else:
        return True

# 始点と終点の関係を検証
# 隣り合っている点同士かどうかを検証する

def validateStartAndEndXY(start_xy, end_xy):
    start_x = start_xy[0]
    start_y = start_xy[1]
    end_x = end_xy[0]
    end_y = end_xy[1]

    start_x_allows = [
        end_x - 1,
        end_x + 1.
    ]
    start_y_allows = [
        end_y - 1,
        end_y + 1,
    ]

    # 隣り合う点同士かどうかを判定
    if (start_x == end_x and start_y in start_y_allows) or (start_y == end_y and start_x in start_x_allows):
        # 隣り合っていたらTrueを戻す
        return True
    else:
        # 隣り合っていなかったら、エラー文を吐き出し、Falseを戻す
        print(Color.RED + '始点と終点は隣り合う点を指定してください' + Color.END)
        return False

# データを更新する

def update(cells, rows, cols, start_xy, end_xy, player, filled):
    if start_xy[0] == end_xy[0]:
        # 始点と終点のX座標が同じであるならば、cols配列を更新
        if cols[start_xy[0]][min(start_xy[1], end_xy[1])] > 0:
            # 既に選択されている場合は、再入力後、再更新する
            print(Color.RED + '既に選択されています。再度選択してください' + Color.END)
            start_xy, end_xy = getInput()
            update(cells, rows, cols, start_xy, end_xy, player, filled)
            return cells, rows, cols, start_xy, end_xy, player, filled
        else:
            cols = updateCols(cols, start_xy, end_xy, player)
    else:
        # そうでなければ、rows配列を更新
        if rows[start_xy[1]][min(start_xy[0], end_xy[0])] > 0:
            # 既に選択されている場合は、再入力後、再更新する
            print(Color.RED + '既に選択されています。再度選択してください' + Color.END)
            start_xy, end_xy = getInput()
            update(cells, rows, cols, start_xy, end_xy, player, filled)
            return cells, rows, cols, start_xy, end_xy, player, filled
        else:
            rows = updateRows(rows, start_xy, end_xy, player)

    #　cells配列を更新
    cells, filled = updateCells(cells, rows, cols, player)

    return cells, rows, cols, start_xy, end_xy, player, filled

# rows配列を更新する

def updateRows(rows, start_xy, end_xy, player):
    # start_xyとend_xyのY座標は同じであることが前提
    # そうでなかった場合は例外を発生させる
    if not start_xy[1] == end_xy[1]:
        raise Exception(Color.RED + '始点と終点のY座標は同じである必要があります。'+ Color.END)

    rows[start_xy[1]][min(start_xy[0], end_xy[0])] = player.number
    return rows

def updateCols(cols, start_xy, end_xy, player):
    # start_xyとend_xyのX座標は同じであることが前提
    # そうでなかった場合は例外を発生させる
    if not start_xy[0] == end_xy[0]:
        raise Exception(Color.RED + '始点と終点のX座標は同じである必要があります。' + Color.END)

    cols[start_xy[0]][min(start_xy[1], end_xy[1])] = player.number
    return cols

def updateCells(cells, rows, cols, player):
    filled = False
    # y: 0-3
    for y, fills_row in enumerate(cells):
        # x: 0-3
        for x, fill in enumerate(fills_row):
            if fill == 0 and rows[y][x] > 0 and rows[y+1][x] > 0 and cols[x][y] > 0 and cols[x+1][y] > 0:
                #　四方に線があるので、マスに現在のターンのプレイヤーの番号を設定する
                cells[y][x] = player.number
                player.count += 1
                filled = True
    return cells, filled

def isEnd(cells):
    return filledCells(cells)

def filledCells(cells):
    # cells配列内の値が全て0以外かどうか
    # すなわち、全てのマスが埋まっているかどうか
    for cell_row in cells:
        for cell in cell_row:
            if cell == 0:
                return False

    return True

def display(cells, rows, cols, players):
    # 水平線
    print(HR)

    for x, dot_row in enumerate(DOTS):
        for y, dot in enumerate(dot_row):
            if y == 0:
                # デザイン調整(左にパディング)
                print(PADDING_LEFT, end='')

            print(dot, end='')

            if y < len(dot_row) - 1:
                if rows[x][y] > 0:
                    print(players[rows[x][y] - 1].color + ROW_LINE + Color.END, end='')
                else:
                    print('　　', end='')

            if y >= len(dot_row) - 1:
                # 2回実行
                for j in range(2):
                    if j == 0:
                        #　改行
                        print('')

                    # デザイン調整(左にパディング)
                    print(PADDING_LEFT, end='')
                    # i:0-3
                    for i in range(len(cols)):
                        if x < len(DOTS) - 1:
                            if cols[i][x] > 0:
                                print(players[cols[i][x] - 1].color +
                                    COL_LINE + Color.END, end='')
                            else:
                                print('　　', end='')

                            if i < len(cols) - 1:
                                if cells[x][i] > 0:
                                    print(
                                        players[cells[x][i] - 1].bg_color + '　　' + Color.END, end='')

                                else:
                                    print('　　', end='')

                    if i >= len(cols) - 1:
                        print('')

    for j, player in enumerate(players):
        # 現在のスコアを表示
        print(player.color + player.name + '：' + str(player.count) + 'マス' + Color.END, end='')
        if j == 0:
            print(' / ', end='')
        if j >= len(players) - 1:
            print('')
            #　水平線
            print(HR)

def yes_no_input(ask, text_color = None):
    if text_color is not None:
        choice = input(text_color + ask + '[y/N]: ' + Color.END)
    else:
        choice = input(ask + '[y/N]: ')

    #小文字に変換
    choice = choice.lower()
    if choice in ['y', 'ye', 'yes']:
        return True
    else:
        return False

if __name__ == '__main__':
    main()


import llm_toys  # For initlializing the hf modules; Look at __init__.py for more details.

import argparse
import random

from peft import get_peft_model, LoraConfig, prepare_model_for_kbit_training
from torch.utils.data import Dataset
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    DataCollatorForLanguageModeling,
    Trainer,
    TrainingArguments,
)

from llm_toys.config import DEFAULT_3B_MODEL, get_bnb_config, TASK_TYPES, TaskType
from llm_toys.utils import dialogue_summary_topic_training_data, paraphrase_tone_training_data

class CustomDataset(Dataset):
    """Custom dataset for the trainer. Tokenizes and keeps the
    entire dataset in memory hence it's not suited for large datasets."""

    def __init__(self, tokenizer: AutoTokenizer, text: str, max_length: int = 1024):
        self.text = text
        self.tokenized = tokenizer(
            self.text,
            padding=True,
            truncation=True,
            max_length=max_length,
            return_tensors="pt",
        )

    def __getitem__(self, ix):
        return self.tokenized["input_ids"][ix]

    def __len__(self):
        return len(self.text)

class PeftQLoraTrainer:
    """Contains helpful methods for loading+transforming the data and training the model."""

    def __init__(self, args: dict):
        self.args = args
        for task_type in args["task_type"].split(","):
            assert (
                task_type.strip().lower() in TASK_TYPES
            ), f"Invalid task type {task_type}; Valid task types are {TASK_TYPES}"
        self.experiment_name = ""
        for task_type in args["task_type"].split(","):
            self.experiment_name += "-".join(task_type.strip().lower().split("_")) + "-"
        self.experiment_name = self.experiment_name[:-1]
        self._models_initialized = False

        if self.args["use_aim"]:
            from aim import Run

            self.run = Run(experiment=self.experiment_name)

    def _set_tokenizer(self) -> None:
        self.tokenizer = AutoTokenizer.from_pretrained(self.args["model_name"])
        self.tokenizer.pad_token = self.tokenizer.eos_token

    def _set_model(self) -> None:
        self.model = AutoModelForCausalLM.from_pretrained(
            self.args["model_name"],
            device_map={"": 0},
            trust_remote_code=True,
            quantization_config=get_bnb_config(),
        )
        self.model.gradient_checkpointing_enable()
        self.model = prepare_model_for_kbit_training(self.model)
        config = LoraConfig(
            r=self.args["lora_r"],
            lora_alpha=self.args["lora_alpha"],
            target_modules=["q_proj", "v_proj"] if "llama-2" in self.args["model_name"].lower() else ["query_key_value"],
            lora_dropout=self.args["lora_dropout"],
            bias="none",
            task_type="CAUSAL_LM",
        )
        self.model = get_peft_model(self.model, config)

    def load_data(self) -> list[str]:
        data = []
        for task_type in self.args["task_type"].split(","):
            task_type = task_type.strip().lower()
            if task_type == TaskType.PARAPHRASE_TONE:
                data += paraphrase_tone_training_data()
            elif task_type == TaskType.DIALOGUE_SUMMARY_TOPIC:
                data += dialogue_summary_topic_training_data()
            else:
                raise ValueError(
                    f"Probably an invalid task type {self.args['task_type']}; Valid task types are {TASK_TYPES}"
                )
        return data

    @staticmethod
    def get_train_evaluation_texts(eval_ratio: float, data: list[str]) -> tuple[list[str], list[str]]:
        val_ixs = set()
        if eval_ratio:
            n_val = int(eval_ratio * len(data))
            val_ixs = set(random.sample(range(len(data)), k=n_val))
        train_ixs = (i for i in range(len(data)) if i not in val_ixs)
        train_texts = [data[i] for i in train_ixs]
        val_texts = [data[i] for i in val_ixs]
        return train_texts, val_texts

    def set_trainer(self, train_dataset: Dataset, eval_dataset: Dataset) -> None:
        training_args = TrainingArguments(
            auto_find_batch_size=True if self.args["batch_size"] else False,
            per_device_train_batch_size=self.args["batch_size"],
            per_device_eval_batch_size=self.args["batch_size"],
            gradient_accumulation_steps=self.args["gradient_accumulation_steps"],
            num_train_epochs=self.args["num_train_epochs"],
            learning_rate=self.args["learning_rate"],
            fp16=True,
            do_eval=self.args["eval_ratio"],
            logging_strategy="steps",
            logging_steps=self.args["logging_steps"],
            evaluation_strategy="steps",
            eval_steps=self.args["eval_steps"],
            output_dir=f"models/{self.args['model_name']}",
            overwrite_output_dir=True,
            save_strategy="epoch",
            optim="paged_adamw_8bit",
            lr_scheduler_type="cosine",
            warmup_ratio=0.05,
        )
        self.trainer = Trainer(
            model=self.model,
            args=training_args,
            data_collator=DataCollatorForLanguageModeling(tokenizer=self.tokenizer, mlm=False),
            train_dataset=train_dataset,
            eval_dataset=eval_dataset,
        )

    def log_params(self):
        """Logs all the metadata and enables a UI to observe/compare them using Aim."""
        hparams = self.args
        self.run["hparams"] = hparams

        for log in self.trainer.state.log_history:
            if "eval_loss" in log:
                self.run.track(log["eval_loss"], name="loss", step=log["step"], context={"subset": "eval"})
            elif "loss" in log:
                self.run.track(log["loss"], name="loss", step=log["step"], context={"subset": "train"})

    def prep_data_train_log_and_save_model(self):
        """Loads the data, preps the train and eval datasets, trains+logs+saves the model."""
        if not self._models_initialized:
            self._set_tokenizer()
            self._set_model()
            self._models_initialized = True

        data = self.load_data()
        print(f"Loaded training data; {len(data)} records")

        random.seed(self.args["seed"])
        random.shuffle(data)
        train_texts, eval_texts = self.get_train_evaluation_texts(self.args["eval_ratio"], data)
        print(f"{'# Train':<15} : {len(train_texts)}")
        if self.args["eval_ratio"]:
            print(f"{'# Eval':<15} : {len(eval_texts)}")

        train_dataset = CustomDataset(self.tokenizer, train_texts, max_length=self.args["max_length"])
        print(f"{'# Train dataset':<15} : {len(train_dataset)}")
        if self.args["eval_ratio"]:
            eval_dataset = CustomDataset(self.tokenizer, eval_texts, max_length=args.max_length)
            print(f"{'# Eval dataset':<15} : {len(eval_dataset)}")

        self.set_trainer(train_dataset, eval_dataset if self.args["eval_ratio"] else None)
        self.model.config.use_cache = False
        self.trainer.train()

        self.model.save_pretrained(f"models/{self.args['model_name']}-{self.experiment_name}")

        if self.args["use_aim"]:
            self.log_params()

    def push_to_hub(self):
        self.model.push_to_hub(f"{self.args['model_name']}-{self.experiment_name}", private=True, use_auth_token=True)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--task_type", type=str, help=f"Valid task types are {TASK_TYPES}.")
    parser.add_argument("--model_name", type=str, default=DEFAULT_3B_MODEL, help="Model name.")
    parser.add_argument("--seed", type=int, default=10, help="Random seed.")
    parser.add_argument("--eval_ratio", type=float, default=0.1, help="Ratio of evaluation data.")
    parser.add_argument("--max_length", type=int, default=512, help="Maximum length of input sequence.")
    parser.add_argument("--lora_r", type=int, default=16, help="Lora r.")
    parser.add_argument("--lora_alpha", type=int, default=32, help="Lora alpha.")
    parser.add_argument("--lora_dropout", type=float, default=0.05, help="Lora dropout.")
    parser.add_argument("--batch_size", type=int, default=1, help="Batch size.")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=4, help="Gradient accumulation steps.")
    parser.add_argument("--num_train_epochs", type=int, default=5, help="# Epochs.")
    parser.add_argument("--learning_rate", type=float, default=2e-4, help="Learning rate.")
    parser.add_argument("--logging_steps", type=int, default=100, help="Logging steps.")
    parser.add_argument("--eval_steps", type=int, default=100, help="Eval steps.")
    parser.add_argument("--use_aim", action=argparse.BooleanOptionalAction, help="To use AIM for logging or not.")
    args = parser.parse_args()

    peft_qlora_trainer = PeftQLoraTrainer(vars(args))
    peft_qlora_trainer.prep_data_train_log_and_save_model()


import torch
import torch.nn as nn
from mmcv.cnn.bricks.registry import TRANSFORMER_LAYER_SEQUENCE
from mmcv.cnn.bricks.transformer import build_transformer_layer
from mmcv.runner.base_module import BaseModule
from projects.mmdet3d_plugin.models.utils.functional import (
    norm_points,
    pos2posemb2d,
    trajectory_coordinate_transform
)

@TRANSFORMER_LAYER_SEQUENCE.register_module()
class MotionTransformerDecoder(BaseModule):
    """Implements the decoder in DETR3D transformer.
    Args:
        return_intermediate (bool): Whether to return intermediate outputs.
        coder_norm_cfg (dict): Config of last normalization layer. Default：
            `LN`.
    """

    def __init__(self, pc_range=None, embed_dims=256, transformerlayers=None, num_layers=3, **kwargs):
        super(MotionTransformerDecoder, self).__init__()
        self.pc_range = pc_range
        self.embed_dims = embed_dims
        self.num_layers = num_layers
        self.intention_interaction_layers = IntentionInteraction()
        self.track_agent_interaction_layers = nn.ModuleList(
            [TrackAgentInteraction() for i in range(self.num_layers)])
        self.map_interaction_layers = nn.ModuleList(
            [MapInteraction() for i in range(self.num_layers)])
        self.bev_interaction_layers = nn.ModuleList(
            [build_transformer_layer(transformerlayers) for i in range(self.num_layers)])

        self.static_dynamic_fuser = nn.Sequential(
            nn.Linear(self.embed_dims*2, self.embed_dims*2),
            nn.ReLU(),
            nn.Linear(self.embed_dims*2, self.embed_dims),
        )
        self.dynamic_embed_fuser = nn.Sequential(
            nn.Linear(self.embed_dims*3, self.embed_dims*2),
            nn.ReLU(),
            nn.Linear(self.embed_dims*2, self.embed_dims),
        )
        self.in_query_fuser = nn.Sequential(
            nn.Linear(self.embed_dims*2, self.embed_dims*2),
            nn.ReLU(),
            nn.Linear(self.embed_dims*2, self.embed_dims),
        )
        self.out_query_fuser = nn.Sequential(
            nn.Linear(self.embed_dims*4, self.embed_dims*2),
            nn.ReLU(),
            nn.Linear(self.embed_dims*2, self.embed_dims),
        )

    def forward(self,
                track_query,
                lane_query,
                track_query_pos=None,
                lane_query_pos=None,
                track_bbox_results=None,
                bev_embed=None,
                reference_trajs=None,
                traj_reg_branches=None,
                agent_level_embedding=None,
                scene_level_ego_embedding=None,
                scene_level_offset_embedding=None,
                learnable_embed=None,
                agent_level_embedding_layer=None,
                scene_level_ego_embedding_layer=None,
                scene_level_offset_embedding_layer=None,
                **kwargs):
        """Forward function for `MotionTransformerDecoder`.
        Args:
            agent_query (B, A, D)
            map_query (B, M, D)
            map_query_pos (B, G, D)
            static_intention_embed (B, A, P, D)
            offset_query_embed (B, A, P, D)
            global_intention_embed (B, A, P, D)
            learnable_intention_embed (B, A, P, D)
            det_query_pos (B, A, D)
        Returns:
            None
        """
        intermediate = []
        intermediate_reference_trajs = []

        B, _, P, D = agent_level_embedding.shape
        track_query_bc = track_query.unsqueeze(2).expand(-1, -1, P, -1)  # (B, A, P, D)
        track_query_pos_bc = track_query_pos.unsqueeze(2).expand(-1, -1, P, -1)  # (B, A, P, D)

        # static intention embedding, which is imutable throughout all layers
        agent_level_embedding = self.intention_interaction_layers(agent_level_embedding)
        static_intention_embed = agent_level_embedding + scene_level_offset_embedding + learnable_embed
        reference_trajs_input = reference_trajs.unsqueeze(4).detach()

        query_embed = torch.zeros_like(static_intention_embed)
        for lid in range(self.num_layers):
            # fuse static and dynamic intention embedding
            # the dynamic intention embedding is the output of the previous layer, which is initialized with anchor embedding
            dynamic_query_embed = self.dynamic_embed_fuser(torch.cat(
                [agent_level_embedding, scene_level_offset_embedding, scene_level_ego_embedding], dim=-1))

            # fuse static and dynamic intention embedding
            query_embed_intention = self.static_dynamic_fuser(torch.cat(
                [static_intention_embed, dynamic_query_embed], dim=-1))  # (B, A, P, D)

            # fuse intention embedding with query embedding
            query_embed = self.in_query_fuser(torch.cat([query_embed, query_embed_intention], dim=-1))

            # interaction between agents
            track_query_embed = self.track_agent_interaction_layers[lid](
                query_embed, track_query, query_pos=track_query_pos_bc, key_pos=track_query_pos)

            # interaction between agents and map
            map_query_embed = self.map_interaction_layers[lid](
                query_embed, lane_query, query_pos=track_query_pos_bc, key_pos=lane_query_pos)

            # interaction between agents and bev, ie. interaction between agents and goals
            # implemented with deformable transformer
            bev_query_embed = self.bev_interaction_layers[lid](
                query_embed,
                value=bev_embed,
                query_pos=track_query_pos_bc,
                bbox_results=track_bbox_results,
                reference_trajs=reference_trajs_input,
                **kwargs)

            # fusing the embeddings from different interaction layers
            query_embed = [track_query_embed, map_query_embed, bev_query_embed, track_query_bc+track_query_pos_bc]
            query_embed = torch.cat(query_embed, dim=-1)
            query_embed = self.out_query_fuser(query_embed)

            if traj_reg_branches is not None:
                # update reference trajectory
                tmp = traj_reg_branches[lid](query_embed)
                bs, n_agent, n_modes, n_steps, _ = reference_trajs.shape
                tmp = tmp.view(bs, n_agent, n_modes, n_steps, -1)

                # we predict speed of trajectory and use cumsum trick to get the trajectory
                tmp[..., :2] = torch.cumsum(tmp[..., :2], dim=3)
                new_reference_trajs = torch.zeros_like(reference_trajs)
                new_reference_trajs = tmp[..., :2]
                reference_trajs = new_reference_trajs.detach()
                reference_trajs_input = reference_trajs.unsqueeze(4)  # BS NUM_AGENT NUM_MODE 12 NUM_LEVEL  2

                # update embedding, which is used in the next layer
                # only update the embedding of the last step, i.e. the goal
                ep_offset_embed = reference_trajs.detach()
                ep_ego_embed = trajectory_coordinate_transform(reference_trajs.unsqueeze(
                    2), track_bbox_results, with_translation_transform=True, with_rotation_transform=False).squeeze(2).detach()
                ep_agent_embed = trajectory_coordinate_transform(reference_trajs.unsqueeze(
                    2), track_bbox_results, with_translation_transform=False, with_rotation_transform=True).squeeze(2).detach()

                agent_level_embedding = agent_level_embedding_layer(pos2posemb2d(
                    norm_points(ep_agent_embed[..., -1, :], self.pc_range)))
                scene_level_ego_embedding = scene_level_ego_embedding_layer(pos2posemb2d(
                    norm_points(ep_ego_embed[..., -1, :], self.pc_range)))
                scene_level_offset_embedding = scene_level_offset_embedding_layer(pos2posemb2d(
                    norm_points(ep_offset_embed[..., -1, :], self.pc_range)))

                intermediate.append(query_embed)
                intermediate_reference_trajs.append(reference_trajs)

        return torch.stack(intermediate), torch.stack(intermediate_reference_trajs)

class TrackAgentInteraction(BaseModule):
    """
    Modeling the interaction between the agents
    """
    def __init__(self,
                 embed_dims=256,
                 num_heads=8,
                 dropout=0.1,
                 batch_first=True,
                 norm_cfg=None,
                 init_cfg=None):
        super().__init__(init_cfg)

        self.batch_first = batch_first
        self.interaction_transformer = nn.TransformerDecoderLayer(d_model=embed_dims,
                                                                  nhead=num_heads,
                                                                  dropout=dropout,
                                                                  dim_feedforward=embed_dims*2,
                                                                  batch_first=batch_first)

    def forward(self, query, key, query_pos=None, key_pos=None):
        '''
        query: context query (B, A, P, D)
        query_pos: mode pos embedding (B, A, P, D)
        key: (B, A, D)
        key_pos: (B, A, D)
        '''
        B, A, P, D = query.shape
        if query_pos is not None:
            query = query + query_pos
        if key_pos is not None:
            key = key + key_pos
        mem = key.expand(B*A, -1, -1)
        # N, A, P, D -> N*A, P, D
        query = torch.flatten(query, start_dim=0, end_dim=1)
        query = self.interaction_transformer(query, mem)
        query = query.view(B, A, P, D)
        return query

class MapInteraction(BaseModule):
    """
    Modeling the interaction between the agent and the map
    """
    def __init__(self,
                 embed_dims=256,
                 num_heads=8,
                 dropout=0.1,
                 batch_first=True,
                 norm_cfg=None,
                 init_cfg=None):
        super().__init__(init_cfg)

        self.batch_first = batch_first
        self.interaction_transformer = nn.TransformerDecoderLayer(d_model=embed_dims,
                                                                  nhead=num_heads,
                                                                  dropout=dropout,
                                                                  dim_feedforward=embed_dims*2,
                                                                  batch_first=batch_first)

    def forward(self, query, key, query_pos=None, key_pos=None):
        '''
        x: context query (B, A, P, D)
        query_pos: mode pos embedding (B, A, P, D)
        '''
        B, A, P, D = query.shape
        if query_pos is not None:
            query = query + query_pos
        if key_pos is not None:
            key = key + key_pos

        # N, A, P, D -> N*A, P, D
        query = torch.flatten(query, start_dim=0, end_dim=1)
        mem = key.expand(B*A, -1, -1)
        query = self.interaction_transformer(query, mem)
        query = query.view(B, A, P, D)
        return query

class IntentionInteraction(BaseModule):
    """
    Modeling the interaction between anchors
    """
    def __init__(self,
                 embed_dims=256,
                 num_heads=8,
                 dropout=0.1,
                 batch_first=True,
                 norm_cfg=None,
                 init_cfg=None):
        super().__init__(init_cfg)

        self.batch_first = batch_first
        self.interaction_transformer = nn.TransformerEncoderLayer(d_model=embed_dims,
                                                                  nhead=num_heads,
                                                                  dropout=dropout,
                                                                  dim_feedforward=embed_dims*2,
                                                                  batch_first=batch_first)

    def forward(self, query):
        B, A, P, D = query.shape
        # B, A, P, D -> B*A,P, D
        rebatch_x = torch.flatten(query, start_dim=0, end_dim=1)
        rebatch_x = self.interaction_transformer(rebatch_x)
        out = rebatch_x.view(B, A, P, D)
        return out


import numpy as np
import scipy.signal as scipysig
import torch
import torch.nn.functional as F
from scipy.stats import chi2
from typing import Tuple, List, NamedTuple
from momentchi2 import hbe

class CollectedData(NamedTuple):
    X: np.ndarray
    U: np.ndarray
    W: np.ndarray
    std_u: float
    std_w: float

class TestStatistics(NamedTuple):
    statistics: float
    pvalue: float

def collect_data(steps: int, std_u: float, std_w: float, sys: scipysig.StateSpace) -> Tuple[np.ndarray, np.ndarray]:
    """ Collect data from a system """
    dim_x, dim_u = sys.B.shape
    U = np.zeros((dim_u, steps))
    X = np.zeros((dim_x, steps + 1))
    W = np.zeros((dim_x, steps))
    X[:, 0] = np.random.normal(size=(dim_x))

    for i in range(steps):
        U[:, i] = std_u * np.random.normal(size=(dim_u))
        W[:, i] = std_w * np.random.normal(size=(dim_x))
        X[:, i+1] = sys.A @ X[:, i] +  np.squeeze(sys.B * U[:, i]) + W[:, i]

    return CollectedData(X, U, W, std_u, std_w)

def correlate(x: np.ndarray, num_lags: int) -> np.ndarray:
    """Correlate a signal"""
    n, T = x.shape
    R = np.zeros((num_lags,n,n))

    for m in range(num_lags):
        for i in range(T-m):
            R[m] += x[:,i:i+1] @ x[:, i+m:i+m+1].T

    return R/T

def torch_correlate(x: torch.Tensor, num_lags: int) -> List[torch.Tensor]:
    """Correlate a signal x using torch"""
    n, T = x.shape
    R = []

    for m in range(num_lags):
        R.append(F.conv1d(x.unsqueeze(1), x[:, m:].unsqueeze(1))[:,:,0]/T)

    return R

def whiteness_test(residuals: np.ndarray, num_lags: int) -> float:
    dim_x, T = residuals.shape
    R = correlate(residuals, num_lags)
    df = (T-dim_x-1) * (dim_x ** 2)
    R0Inv = np.linalg.inv(R[0])
    statistics = T * np.sum([np.trace(R[x].T @ R0Inv @ R[x] @ R0Inv) for x in range(1, num_lags)])
    q = chi2.cdf(statistics, df)
    return statistics

def pvalue_whiteness_test(R: np.ndarray, num_lags: int, T: int) -> TestStatistics:
    """Run whiteness test on the residuals

    Args:
        R (np.ndarray): correlation matrices, should be of shape (L,n,n), where L is the number of lags
        num_lags (int): number of lags to use to run the test. Should be lower than L
        T (int): Horizon length

    Returns:
        TestStatistics: Return the statistics of the test
    """
    R_lags, dim_x, _ = R.shape
    assert R_lags > num_lags, f"R has been computed only for {R_lags} lags, not for {num_lags} lags."
    df = (T-dim_x-1) * (dim_x ** 2)
    R0Inv = np.linalg.inv(R[0])
    statistics = T * np.sum([np.trace(R[x].T @ R0Inv @ R[x] @ R0Inv) for x in range(1, num_lags + 1)])
    p = chi2.cdf(statistics, df)
    return TestStatistics(statistics, p)

def pvalue_residuals_variance_test(residuals: np.ndarray, dim_u: int, sigma_eigs: List[float]) -> TestStatistics:
    """
    Compute the pvalue under H0 that the residuals follow a linear combination of chi_squared distribution

    :param residuals: a (dim_x, T) matrix with the residuals
    :param dim_u: dimension of the input
    :param sigma_eigs: list of eigenvalues of the noise covariance matrix
    :return pval: pvalue
    """
    assert isinstance(sigma_eigs, list), "Sigma_eigs needs to be a list"
    dim_x, T = residuals.shape
    assert len(sigma_eigs) == dim_x, "Sigma_eigs needs to have dim_x values"
    eigs = sigma_eigs * (T-dim_x-dim_u)
    fro_norm_squared  = np.linalg.norm(residuals, 'fro' ) ** 2
    p = hbe(coeff=eigs, x=fro_norm_squared)
    return TestStatistics(fro_norm_squared, 2*min(p, 1-p))


from typing import List

import torch

from roar.core.classes import Loss, typecheck
from roar.core.neural_types import LossType, NeuralType

__all__ = ["AggregatorLoss"]

class AggregatorLoss(Loss):
    """
    Sums several losses into one.

    Args:
        num_inputs: number of input losses
        weights: a list of coefficient for merging losses
    """

    @property
    def input_types(self):
        """Returns definitions of module input ports."""
        input_types = {}
        for i in range(self._num_losses):
            input_types["loss_" + str(i + 1)] = NeuralType(elements_type=LossType())

        return input_types

    @property
    def output_types(self):
        """Returns definitions of module output ports."""
        return {"loss": NeuralType(elements_type=LossType())}

    def __init__(self, num_inputs: int = 2, weights: List[float] = None):
        super().__init__()
        self._num_losses = num_inputs
        if weights is not None and len(weights) != num_inputs:
            raise ValueError(
                "Length of weights should be equal to the number of inputs (num_inputs)"
            )

        self._weights = weights

    @typecheck()
    def forward(self, **kwargs):
        values = [kwargs[x] for x in sorted(kwargs.keys())]
        loss = torch.zeros_like(values[0])
        for loss_idx, loss_value in enumerate(values):
            if self._weights is not None:
                loss = loss.add(loss_value, alpha=self._weights[loss_idx])
            else:
                loss = loss.add(loss_value)
        return loss