import os
import gc
import json
import torch
import shutil
import operator
import wandb
import numpy as np

from torch import nn
from pathlib import Path
from itertools import chain
from transformers import AutoModel
from torch.optim.lr_scheduler import CyclicLR
from sklearn.metrics import mean_squared_error
from torch.utils.data import Dataset, DataLoader
from transformers import get_cosine_schedule_with_warmup, get_linear_schedule_with_warmup, get_cosine_with_hard_restarts_schedule_with_warmup, get_polynomial_decay_schedule_with_warmup

def clear_cuda():
    gc.collect()
    torch.cuda.empty_cache()

class UnoTextDataset(Dataset):
    def __init__(self, text_excerpts, targets):
        self.text_excerpts = text_excerpts
        self.targets = targets

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

    def __getitem__(self, idx):
        sample = {'text_excerpt': self.text_excerpts[idx],
                  'target': self.targets[idx]}
        return sample

class DuoTextDataset(Dataset):
    """
    If the first input is easier to read then the target is 1
    If the second input is easier to read then the target is -1
    """
    def __init__(self, text_excerpts_left, text_excerpts_right, targets):
        self.text_excerpts_left = text_excerpts_left
        self.text_excerpts_right = text_excerpts_right
        self.targets = targets

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

    def __getitem__(self, idx):
        sample = {'text_excerpt_left': self.text_excerpts_left[idx],
                  'text_excerpt_right': self.text_excerpts_right[idx],
                  'target': self.targets[idx]}
        return sample

def create_uno_text_dataloader(data, batch_size, shuffle, sampler, apply_preprocessing=True, num_workers=4, pin_memory=True, drop_last=False):
    # Preprocessing
    if apply_preprocessing:
        data['excerpt'] = data['excerpt'].apply(lambda x: x.replace('\n', ' '))
        data['excerpt'] = data['excerpt'].apply(lambda x: ' '.join(x.split()))

    text_excerpts = data['excerpt'].tolist()
    targets = data['target'].to_numpy().astype(np.float32).reshape(-1, 1)
    dataset = UnoTextDataset(text_excerpts=text_excerpts, targets=targets)
    dataloader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=shuffle, sampler=sampler,
                            num_workers=num_workers, pin_memory=pin_memory, drop_last=drop_last)
    return dataloader

def create_duo_text_dataloader(data, batch_size, shuffle, sampler, apply_preprocessing=True, num_workers=4, pin_memory=True, drop_last=False):
    if apply_preprocessing:
        data['easy_text'] = data['easy_text'].apply(lambda x: x.replace('\n', ' '))
        data['easy_text'] = data['easy_text'].apply(lambda x: ' '.join(x.split()))

        data['difficult_text'] = data['difficult_text'].apply(lambda x: x.replace('\n', ' '))
        data['difficult_text'] = data['difficult_text'].apply(lambda x: ' '.join(x.split()))

    text_excerpts_left = data['easy_text'].tolist()
    text_excerpts_right = data['difficult_text'].tolist()
    targets = len(data) * [1]
    targets = np.asarray(targets).astype(np.float32).reshape(-1, 1)

    dataset = DuoTextDataset(text_excerpts_left=text_excerpts_left,
                             text_excerpts_right=text_excerpts_right,
                             targets=targets)
    dataloader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=shuffle, sampler=sampler,
                            num_workers=num_workers, pin_memory=pin_memory, drop_last=drop_last)
    return dataloader

class AttentionHead(nn.Module):
    def __init__(self, input_dim, head_hidden_dim):
        super(AttentionHead, self).__init__()
        head_hidden_dim = input_dim if head_hidden_dim is None else head_hidden_dim
        self.W = nn.Linear(input_dim, head_hidden_dim)
        self.V = nn.Linear(head_hidden_dim, 1)

    def forward(self, x):
        attention_scores = self.V(torch.tanh(self.W(x)))
        attention_scores = torch.softmax(attention_scores, dim=1)
        attentive_x = attention_scores * x
        attentive_x = attentive_x.sum(axis=1)
        return attentive_x

class MaskFilledAttentionHead(nn.Module):
    def __init__(self, input_dim, head_hidden_dim):
        super(MaskFilledAttentionHead, self).__init__()
        head_hidden_dim = input_dim if head_hidden_dim is None else head_hidden_dim
        self.W = nn.Linear(input_dim, head_hidden_dim)
        self.V = nn.Linear(head_hidden_dim, 1)

    def forward(self, x, attention_mask):
        attention_scores = self.V(torch.tanh(self.W(x)))
        attention_scores[attention_mask==0] = -10
        attention_scores = torch.softmax(attention_scores, dim=1)
        attentive_x = attention_scores * x
        attentive_x = attentive_x.sum(axis=1)
        return attentive_x

class MaskAddedAttentionHead(nn.Module):
    def __init__(self, input_dim, head_hidden_dim):
        super(MaskAddedAttentionHead, self).__init__()
        head_hidden_dim = input_dim if head_hidden_dim is None else head_hidden_dim
        self.W = nn.Linear(input_dim, head_hidden_dim)
        self.V = nn.Linear(head_hidden_dim, 1)

    def forward(self, x, attention_mask):
        attention_scores = self.V(torch.tanh(self.W(x)))
        attention_scores = attention_scores + attention_mask
        attention_scores = torch.softmax(attention_scores, dim=1)
        attentive_x = attention_scores * x
        attentive_x = attentive_x.sum(axis=1)
        return attentive_x

class RobertaPoolerOutputRegressor(nn.Module):
    def __init__(self, model_path, dropout_prob=0.1,  roberta_hidden_dropout_prob=0.1, roberta_attention_probs_dropout_prob=0.1, **kwargs):
        super(RobertaPoolerOutputRegressor, self).__init__()
        self.roberta = AutoModel.from_pretrained(model_path,
                                                 hidden_dropout_prob=roberta_hidden_dropout_prob,
                                                 attention_probs_dropout_prob=roberta_attention_probs_dropout_prob, **kwargs)
        self.dropout = nn.Dropout(dropout_prob)
        self.regressor = nn.Linear(self.roberta.config.hidden_size, 1)

    def forward(self, inputs):
        roberta_outputs = self.roberta(**inputs)
        pooler_output = roberta_outputs['pooler_output']
        pooler_output = self.dropout(pooler_output)
        logits = self.regressor(pooler_output)
        return logits

class RobertaLastHiddenStateRegressor(nn.Module):
    def __init__(self, model_path, head_hidden_dim=None, dropout_prob=0.1, roberta_hidden_dropout_prob=0.1, roberta_attention_probs_dropout_prob=0.1, **kwargs):
        super(RobertaLastHiddenStateRegressor, self).__init__()
        self.roberta = AutoModel.from_pretrained(model_path,
                                                 hidden_dropout_prob=roberta_hidden_dropout_prob,
                                                 attention_probs_dropout_prob=roberta_attention_probs_dropout_prob, **kwargs)
        self.head = AttentionHead(input_dim=self.roberta.config.hidden_size, head_hidden_dim=head_hidden_dim)
        self.dropout = nn.Dropout(dropout_prob)
        self.regressor = nn.Linear(self.roberta.config.hidden_size, 1)

    def forward(self, inputs):
        roberta_outputs = self.roberta(**inputs)
        last_hidden_state = roberta_outputs['last_hidden_state']
        attentive_vector = self.head(last_hidden_state)
        attentive_vector = self.dropout(attentive_vector)
        logits = self.regressor(attentive_vector)
        return logits

class RobertaMaskedLastHiddenStateRegressor(nn.Module):
    def __init__(self, model_path, head_hidden_dim=None, dropout_prob=0.1, roberta_hidden_dropout_prob=0.1, roberta_attention_probs_dropout_prob=0.1, **kwargs):
        super(RobertaMaskedLastHiddenStateRegressor, self).__init__()
        self.roberta = AutoModel.from_pretrained(model_path,
                                                 hidden_dropout_prob=roberta_hidden_dropout_prob,
                                                 attention_probs_dropout_prob=roberta_attention_probs_dropout_prob, **kwargs)
        self.head = AttentionHead(input_dim=self.roberta.config.hidden_size, head_hidden_dim=head_hidden_dim)
        self.dropout = nn.Dropout(dropout_prob)
        self.regressor = nn.Linear(self.roberta.config.hidden_size, 1)

    def forward(self, inputs):
        roberta_outputs = self.roberta(**inputs)
        last_hidden_state = roberta_outputs['last_hidden_state']
        masked_last_hidden_state = last_hidden_state * torch.unsqueeze(inputs['attention_mask'], dim=2)
        attentive_vector = self.head(masked_last_hidden_state)
        attentive_vector = self.dropout(attentive_vector)
        logits = self.regressor(attentive_vector)
        return logits

class RobertaMaskFilledAttentionHeadRegressor(nn.Module):
    def __init__(self, model_path, dropout_prob=0.1, head_hidden_dim=None, roberta_hidden_dropout_prob=0.1, roberta_attention_probs_dropout_prob=0.1, **kwargs):
        super(RobertaMaskFilledAttentionHeadRegressor, self).__init__()
        self.roberta = AutoModel.from_pretrained(model_path,
                                                 hidden_dropout_prob=roberta_hidden_dropout_prob,
                                                 attention_probs_dropout_prob=roberta_attention_probs_dropout_prob, **kwargs)
        self.head = MaskFilledAttentionHead(input_dim=self.roberta.config.hidden_size, head_hidden_dim=head_hidden_dim)
        self.dropout = nn.Dropout(dropout_prob)
        self.regressor = nn.Linear(self.roberta.config.hidden_size, 1)

    def forward(self, inputs):
        roberta_outputs = self.roberta(**inputs)
        last_hidden_state = roberta_outputs['last_hidden_state']
        attentive_vector = self.head(last_hidden_state, torch.unsqueeze(inputs['attention_mask'], dim=2))
        attentive_vector = self.dropout(attentive_vector)
        logits = self.regressor(attentive_vector)
        return logits

class RobertaMaskAddedAttentionHeadRegressor(nn.Module):
    def __init__(self, model_path, dropout_prob=0.1, head_hidden_dim=None, roberta_hidden_dropout_prob=0.1, roberta_attention_probs_dropout_prob=0.1, **kwargs):
        super(RobertaMaskAddedAttentionHeadRegressor, self).__init__()
        self.roberta = AutoModel.from_pretrained(model_path,
                                                 hidden_dropout_prob=roberta_hidden_dropout_prob,
                                                 attention_probs_dropout_prob=roberta_attention_probs_dropout_prob, **kwargs)
        self.head = MaskAddedAttentionHead(input_dim=self.roberta.config.hidden_size, head_hidden_dim=head_hidden_dim)
        self.dropout = nn.Dropout(dropout_prob)
        self.regressor = nn.Linear(self.roberta.config.hidden_size, 1)

    def forward(self, inputs):
        roberta_outputs = self.roberta(**inputs)
        last_hidden_state = roberta_outputs['last_hidden_state']
        attentive_vector = self.head(last_hidden_state, torch.unsqueeze(inputs['attention_mask'], dim=2))
        attentive_vector = self.dropout(attentive_vector)
        logits = self.regressor(attentive_vector)
        return logits

class RobertaNHiddenStateRegressor(nn.Module):
    def __init__(self, model_path, dropout_prob=0.1, head_hidden_dim=None, roberta_hidden_dropout_prob=0.1, roberta_attention_probs_dropout_prob=0.1, **kwargs):
        super(RobertaNHiddenStateRegressor, self).__init__()
        self.num_last_hidden_states = kwargs.pop('num_last_hidden_states')
        self.roberta = AutoModel.from_pretrained(model_path,
                                                 hidden_dropout_prob=roberta_hidden_dropout_prob,
                                                 attention_probs_dropout_prob=roberta_attention_probs_dropout_prob, **kwargs)
        self.head = MaskAddedAttentionHead(input_dim=self.roberta.config.hidden_size * self.num_last_hidden_states, head_hidden_dim=head_hidden_dim)
        self.dropout = nn.Dropout(dropout_prob)
        self.regressor = nn.Linear(self.roberta.config.hidden_size *  self.num_last_hidden_states, 1)

    def forward(self, inputs):
        roberta_outputs = self.roberta(**inputs, output_hidden_states=True)
        last_hidden_states = torch.cat((roberta_outputs['hidden_states'][- self.num_last_hidden_states:]), axis=2)
        attentive_vector = self.head(last_hidden_states, torch.unsqueeze(inputs['attention_mask'], dim=2))
        attentive_vector = self.dropout(attentive_vector)
        logits = self.regressor(attentive_vector)
        return logits

class RobertaLastHiddenStateMeanPooler(nn.Module):
    def __init__(self, model_path, dropout_prob=0.1, head_hidden_dim=None,  roberta_hidden_dropout_prob=0.1, roberta_attention_probs_dropout_prob=0.1, **kwargs):
        super(RobertaLastHiddenStateMeanPooler, self).__init__()
        self.roberta = AutoModel.from_pretrained(model_path,
                                                 hidden_dropout_prob=roberta_hidden_dropout_prob,
                                                 attention_probs_dropout_prob=roberta_attention_probs_dropout_prob, **kwargs)
        self.dropout = nn.Dropout(dropout_prob)
        self.regressor = nn.Linear(self.roberta.config.hidden_size, 1)

    def forward(self, inputs):
        roberta_outputs = self.roberta(**inputs)
        last_hidden_state = roberta_outputs['last_hidden_state']
        masked_last_hidden_state = last_hidden_state * torch.unsqueeze(inputs['attention_mask'], dim=2)
        num_tokens = torch.unsqueeze(inputs['attention_mask'], dim=2)
        num_tokens = torch.clamp(num_tokens, min=1e-9)
        mean_embeddings = masked_last_hidden_state.sum(axis=1) / num_tokens.sum(axis=1)
        mean_embeddings = self.dropout(mean_embeddings)
        logits = self.regressor(mean_embeddings)
        return logits

def split_into_wd_groups(param_group, weight_decay):
    # Applies weight decay
    weight_parameters = {'params': [param_group['params'][index] for index, name in enumerate(param_group['param_names']) if 'weight' in name and 'LayerNorm' not in name],
                         'param_names': [param_group['param_names'][index] for index, name in enumerate(param_group['param_names']) if 'weight' in name and 'LayerNorm' not in name],
                         'lr': param_group['lr'],
                         'weight_decay': weight_decay,
                         'name': param_group['name']+'_weight'}
    # Does not apply weight decay
    bias_ln_parameters = {'params': [param_group['params'][index] for index, name in enumerate(param_group['param_names']) if 'bias' in name or 'LayerNorm' in name],
                          'param_names': [param_group['param_names'][index] for index, name in enumerate(param_group['param_names']) if 'bias' in name or 'LayerNorm' in name],
                          'lr': param_group['lr'],
                          'weight_decay': 0.0,
                          'name': param_group['name']+'_bias_ln'}
    parameters = [weight_parameters, bias_ln_parameters]
    return parameters

def get_optimizer_parameters(group_mode, lr, model, **kwargs):
    non_bert = [(n,p) for n,p in model.named_parameters() if 'roberta' not in n]
    no_decay = ['bias', 'gamma', 'beta']

    if group_mode == 'i':
        optimizer_parameters = [{'params': [p for n,p in model.named_parameters() if ('bias' not in n) and ('pooler' not in n)],
                                 'lr': lr, 'weight_decay':0.01, 'name': 'weights'},
                                {'params': [p for n,p in model.named_parameters() if ('bias' in n) and ('pooler' not in n)],
                                 'lr': lr, 'weight_decay':0.00, 'name': 'bias'}]

    elif group_mode == 'j':
        optimizer_parameters = [{'params': [p for n,p in model.named_parameters() if ('bias' not in n)],
                                 'lr': lr, 'weight_decay':0.01, 'name': 'weights'},
                                {'params': [p for n,p in model.named_parameters() if ('bias' in n)],
                                 'lr': lr, 'weight_decay':0.00, 'name': 'bias'}]

    elif group_mode == 'k':
        # Finetuning task specific layers
        optimizer_parameters = [{'params': [p for n,p in non_bert if 'bias' not in n],
                                 'param_names': [n for n,p in non_bert if 'bias' not in n],
                                 'lr': lr, 'weight_decay':0.01, 'name': 'non_roberta_weights'},
                                {'params': [p for n,p in non_bert if 'bias' in n],
                                 'param_names': [n for n,p in non_bert if 'bias' in n],
                                 'lr': lr, 'weight_decay':0.00, 'name': 'non_roberta_bias'}]

    elif group_mode  == 's':
        multiplicative_factor = kwargs['multiplicative_factor']
        optimizer_parameters = [{'params': [p for n,p in model.roberta.named_parameters() if (not any(nd in n for nd in no_decay)) and ('pooler' not in n)],
                                 'lr': lr, 'weight_decay' : 0.01, 'name': 'roberta_weights'},
                                {'params': [p for n,p in model.roberta.named_parameters() if (any(nd in n for nd in no_decay))  and ('pooler' not in n)],
                                 'lr': lr,'weight_decay': 0.0, 'name': 'roberta_bias'},
                                {'params': [p for n,p in model.named_parameters() if all(nd not in n for nd in ['roberta','bias'])],
                                 'lr': lr * multiplicative_factor, 'weight_decay':0.01, 'name': 'non_roberta_weights'},
                                {'params': [p for n,p in non_bert if 'bias' in n],
                                 'lr': lr * multiplicative_factor, 'weight_decay':0.00, 'name': 'non_roberta_bias'}]
    elif group_mode == 'b':
        multiplicative_factor = kwargs['multiplicative_factor']
        model_parameters = {name: param for name, param in model.named_parameters()}
        model_parameters_names = [name for name, param in model.named_parameters()]

        # LR for task specific layers
        if kwargs['train_pooler']:
            task_specific_layer_names = [name for name, param in model.named_parameters() if 'regressor' in name or 'head' in name or 'pooler' in name or 'layer_norm' in name]
        else:
            task_specific_layer_names = [name for name, param in model.named_parameters() if 'regressor' in name or 'head' in name or 'layer_norm' in name]
        task_specific_layer_params = [model_parameters.get(name) for name in task_specific_layer_names]
        task_specific_optimizer_parameters = [{'params': task_specific_layer_params,
                                               'param_names': task_specific_layer_names,
                                               'lr': lr,
                                               'name': 'task_specific_layers'}]
        # LR for roberta layers
        # Freeze embeddings
        roberta_layer_names = [name for name, param in model.named_parameters() if 'roberta' in name]
        max_num_layers = model.roberta.config.num_hidden_layers
        roberta_layers_groups = {layer_num: {'params': [],
                                             'param_names': [],
                                             'lr': lr * multiplicative_factor ** (max_num_layers - layer_num),
                                             'name': f'layer_{layer_num}'} for layer_num in range(max_num_layers)}
        for layer_num in range(max_num_layers):
            for layer_name in roberta_layer_names:
                if f'layer.{layer_num}.' in layer_name:
                    roberta_layers_groups[layer_num]['param_names'].append(layer_name)
                    roberta_layers_groups[layer_num]['params'].append(model_parameters.get(layer_name))
        roberta_layers_optimizer_parameters = list(roberta_layers_groups.values())
        # Combine task specific layers and roberta layers
        optimizer_parameters = roberta_layers_optimizer_parameters + task_specific_optimizer_parameters

    elif group_mode == 'b_wd':
        multiplicative_factor = kwargs['multiplicative_factor']
        model_parameters = {name: param for name, param in model.named_parameters()}
        model_parameters_names = [name for name, param in model.named_parameters()]

        # LR for task specific layers
        if kwargs['train_pooler']:
            task_specific_layer_names = [name for name, param in model.named_parameters() if 'regressor' in name or 'head' in name or 'pooler' in name or 'layer_norm' in name]
        else:
            task_specific_layer_names = [name for name, param in model.named_parameters() if 'regressor' in name or 'head' in name or 'layer_norm' in name]
        task_specific_layer_params = [model_parameters.get(name) for name in task_specific_layer_names]
        task_specific_optimizer_parameters = [{'params': task_specific_layer_params,
                                               'param_names': task_specific_layer_names,
                                               'lr': lr,
                                               'name': 'task_specific_layers'}]
        # LR for roberta layers
        # Freeze embeddings
        roberta_layer_names = [name for name, param in model.named_parameters() if 'roberta' in name]
        max_num_layers = model.roberta.config.num_hidden_layers
        roberta_layers_groups = {layer_num: {'params': [],
                                             'param_names': [],
                                             'lr': lr * multiplicative_factor ** (max_num_layers - layer_num),
                                             'name': f'layer_{layer_num}'} for layer_num in range(max_num_layers)}
        for layer_num in range(max_num_layers):
            for layer_name in roberta_layer_names:
                if f'layer.{layer_num}.' in layer_name:
                    roberta_layers_groups[layer_num]['param_names'].append(layer_name)
                    roberta_layers_groups[layer_num]['params'].append(model_parameters.get(layer_name))
        roberta_layers_optimizer_parameters = list(roberta_layers_groups.values())
        # Combine task specific layers and roberta layers
        optimizer_parameters_without_wd = roberta_layers_optimizer_parameters + task_specific_optimizer_parameters

        # Assign weight decay
        weight_decay = kwargs['weight_decay']
        optimizer_parameters = []
        for layer_parameters in optimizer_parameters_without_wd:
            weight_parameters = {'params': [], 'param_names': [], 'lr': layer_parameters['lr'], 'name': layer_parameters['name']+'_weights', 'weight_decay': weight_decay}
            bias_parameters = {'params': [], 'param_names': [], 'lr': layer_parameters['lr'], 'name': layer_parameters['name']+'_bias', 'weight_decay': 0.0}
            layer_norm_parameters = {'params': [], 'param_names': [], 'lr': layer_parameters['lr'], 'name': layer_parameters['name']+'_layer_norm', 'weight_decay': 0.0}
            for param, param_name in zip(layer_parameters['params'], layer_parameters['param_names']):
                if 'LayerNorm' in param_name:
                    layer_norm_parameters['params'].append(param)
                    layer_norm_parameters['param_names'].append(param_name)
                elif 'bias' in param_name:
                    bias_parameters['params'].append(param)
                    bias_parameters['param_names'].append(param_name)
                else:
                    weight_parameters['params'].append(param)
                    weight_parameters['param_names'].append(param_name)
            optimizer_parameters.append(weight_parameters)
            optimizer_parameters.append(bias_parameters)
            optimizer_parameters.append(layer_norm_parameters)

    elif group_mode == 'be_wd':
        multiplicative_factor = kwargs['multiplicative_factor']
        weight_decay = kwargs['weight_decay']
        max_num_layers = model.roberta.config.num_hidden_layers

        # Task Specific Layer group
        tsl_param_group = [{'params': [param for name, param in model.named_parameters() if 'roberta' not in name],
                            'param_names': [name for name, param in model.named_parameters() if 'roberta' not in name],
                            'lr': lr,
                            'name': 'tsl'}]

        # Roberta Layer group
        roberta_layers_param_groups = []
        for layer_num in reversed(range(max_num_layers)):
            roberta_layer_param_groups = {'params': [param for name, param in model.named_parameters() if f'roberta.encoder.layer.{layer_num}.' in name],
                                          'param_names': [name for name, param in model.named_parameters() if f'roberta.encoder.layer.{layer_num}.' in name],
                                          'lr': lr * (multiplicative_factor ** (max_num_layers - layer_num)),
                                          'name': f'layer_{layer_num}'}
            roberta_layers_param_groups.append(roberta_layer_param_groups)

        # Embeddding group
        embedding_lr = lr * (multiplicative_factor ** (max_num_layers + 1))
        embedding_param_group = [{'params': [param for name, param in model.named_parameters() if 'embedding' in name],
                                  'param_names': [name for name, param in model.named_parameters() if 'embedding' in name],
                                  'lr': embedding_lr,
                                  'name': 'embedding'}]

        param_groups = tsl_param_group + roberta_layers_param_groups + embedding_param_group
        optimizer_parameters = list(chain(*[split_into_wd_groups(param_group, weight_decay=weight_decay) for param_group in param_groups]))

    elif group_mode == 'a':
        group1 = ['layer.0.','layer.1.','layer.2.','layer.3.']
        group2 = ['layer.4.','layer.5.','layer.6.','layer.7.']
        group3 = ['layer.8.','layer.9.','layer.10.','layer.11.']
        group_all = group1 + group2 + group3
        optimizer_parameters = [
            {'params': [p for n, p in model.roberta.named_parameters() if not any(nd in n for nd in no_decay) and any(nd in n for nd in group1)],'weight_decay_rate': 0.01, 'lr': lr/2.6, 'name': 'roberta_group_1_weights'},
            {'params': [p for n, p in model.roberta.named_parameters() if not any(nd in n for nd in no_decay) and any(nd in n for nd in group2)],'weight_decay_rate': 0.01, 'lr': lr, 'name': 'roberta_group_2_weights'},
            {'params': [p for n, p in model.roberta.named_parameters() if not any(nd in n for nd in no_decay) and any(nd in n for nd in group3)],'weight_decay_rate': 0.01, 'lr': lr*2.6, 'name': 'roberta_group_3_weights'},
            {'params': [p for n, p in model.roberta.named_parameters() if any(nd in n for nd in no_decay) and any(nd in n for nd in group1)],'weight_decay_rate': 0.0, 'lr': lr/2.6, 'name': 'roberta_group_1_bias'},
            {'params': [p for n, p in model.roberta.named_parameters() if any(nd in n for nd in no_decay) and any(nd in n for nd in group2)],'weight_decay_rate': 0.0, 'lr': lr, 'name': 'roberta_group_2_bias'},
            {'params': [p for n, p in model.roberta.named_parameters() if any(nd in n for nd in no_decay) and any(nd in n for nd in group3)],'weight_decay_rate': 0.0, 'lr': lr*2.6, 'name': 'roberta_group_3_bias'},

            {'params': [p for n, p in non_bert if  'bias' not in n and 'head' in n], 'lr':lr*10, "momentum" : 0.99,'weight_decay_rate':0.01, 'name': 'head_weights'},
            {'params': [p for n, p in non_bert if  'bias' in n and 'head' in n], 'lr':lr*10, "momentum" : 0.99,'weight_decay_rate':0.0, 'name': 'head_bias'},

            {'params': [p for n, p in non_bert if  'bias' not in n and 'regressor' in n], 'lr':lr*10, "momentum" : 0.99,'weight_decay_rate':0.01, 'name': 'regressor_weights'},
            {'params': [p for n, p in non_bert if "bias"  in n and 'regressor' in n], 'lr':lr*10, "momentum" : 0.99,'weight_decay_rate':0.00, 'name': 'regressor_bias'},
        ]
    return optimizer_parameters

def make_cyclic_scheduler(optimizer, **kwargs):
    base_lr = kwargs['base_lr']
    step_size_up = kwargs['step_size_up']
    max_lr = kwargs['max_lr']
    cycle_momentum = kwargs['cycle_momentum']
    scheduler = CyclicLR(optimizer=optimizer, base_lr=base_lr, max_lr=max_lr, step_size_up=step_size_up, cycle_momentum=cycle_momentum)
    return scheduler

def make_cosine_schedule_with_warmup(optimizer, **kwargs):
    num_warmup_steps = kwargs['num_warmup_steps']
    num_training_steps = kwargs['num_training_steps']
    scheduler = get_cosine_schedule_with_warmup(optimizer=optimizer,
                                                num_warmup_steps=num_warmup_steps,
                                                num_training_steps=num_training_steps)
    return scheduler

def make_linear_schedule_with_warmup(optimizer, **kwargs):
    num_warmup_steps = kwargs['num_warmup_steps']
    num_training_steps = kwargs['num_training_steps']
    scheduler = get_linear_schedule_with_warmup(optimizer=optimizer,
                                                num_warmup_steps=num_warmup_steps,
                                                num_training_steps=num_training_steps)
    return scheduler

def make_cosine_with_hard_restarts_schedule_with_warmup(optimizer, **kwargs):
    num_warmup_steps = kwargs['num_warmup_steps']
    num_training_steps = kwargs['num_training_steps']
    scheduler = get_cosine_with_hard_restarts_schedule_with_warmup(optimizer=optimizer,
                                                                   num_warmup_steps=num_warmup_steps,
                                                                   num_training_steps=num_training_steps,
                                                                   num_cycles=3)
    return scheduler

def make_polynomial_decay_schedule_with_warmup(optimizer, **kwargs):
    num_warmup_steps = kwargs['num_warmup_steps']
    num_training_steps = kwargs['num_training_steps']
    scheduler = get_polynomial_decay_schedule_with_warmup(optimizer=optimizer,
                                                          num_warmup_steps=num_warmup_steps,
                                                          num_training_steps=num_training_steps,
                                                          power=3.0)
    return scheduler

def get_scheduler(scheduler_type, optimizer, **kwargs):
    if scheduler_type == 'cyclic':
        scheduler = make_cyclic_scheduler(optimizer, **kwargs)
    elif scheduler_type == 'cosine_schedule_with_warmup':
        scheduler = make_cosine_schedule_with_warmup(optimizer, **kwargs)
    elif scheduler_type == 'get_linear_schedule_with_warmup':
        scheduler = make_linear_schedule_with_warmup(optimizer, **kwargs)
    elif scheduler_type == 'cosine_with_hard_restarts_schedule_with_warmup':
        scheduler = make_cosine_with_hard_restarts_schedule_with_warmup(optimizer, **kwargs)
    elif scheduler_type == 'polynomial_decay_schedule_with_warmup':
        scheduler = make_polynomial_decay_schedule_with_warmup(optimizer, **kwargs)
    else:
        scheduler = None
    return scheduler

def compute_predictions(text_excerpts, tokenizer, model, max_length, device, **kwargs):
    if max_length is None:
        # Sequence bucketing
        inputs = tokenizer(text=text_excerpts, padding=True, truncation=True, return_tensors='pt')
    else:
        inputs = tokenizer(text=text_excerpts, padding='max_length', truncation=True, max_length=max_length,  return_tensors='pt')
    inputs = {key:value.to(device) for key, value in inputs.items()}
    predictions = model(inputs)
    return predictions

def forward_pass_uno_text_batch(batch, tokenizer, model, compute_loss_fn, max_length, device, **kwargs):
    predictions = compute_predictions(text_excerpts=batch['text_excerpt'], tokenizer=tokenizer, model=model, max_length=max_length, device=device, **kwargs)
    outputs = {'predictions': predictions}
    loss = compute_loss_fn(outputs=outputs, targets=batch['target'], device=device, **kwargs) if compute_loss_fn is not None else None
    outputs['loss'] = loss
    return outputs

mse_loss_fn = nn.MSELoss()

def compute_mse_loss(outputs, targets, device, **kwargs):
    predictions = outputs['predictions']
    targets = targets.to(device)
    loss = mse_loss_fn(predictions, targets)
    return loss

def compute_rmse_loss(outputs, targets, device, **kwargs):
    predictions = outputs['predictions']
    targets = targets.to(device)
    loss = torch.sqrt(mse_loss_fn(predictions, targets))
    return loss

ranking_loss_fn = nn.MarginRankingLoss()

def compute_ranking_loss(outputs, targets, device, **kwargs):
    predictions_left = outputs['predictions_left']
    predictions_right = outputs['predictions_right']
    targets = targets.to(device)

    predictions_left = predictions_left.reshape(len(predictions_left))
    predictions_right = predictions_right.reshape(len(predictions_right))
    targets = targets.reshape(len(targets))

    loss = ranking_loss_fn(predictions_left, predictions_right, targets)
    return loss

def compute_rmse_score(outputs, targets, **kwargs):
    predictions = outputs['predictions']
    predictions = predictions.detach().cpu().numpy()
    targets = targets.detach().cpu().numpy()
    rmse_score = mean_squared_error(targets, predictions, squared=False)
    return rmse_score

class UnoStacker:
    def __init__(self):
        self.predictions = []
        self.targets = []

    def update(self, batch_outputs, batch_targets):
        self.predictions.append(batch_outputs['predictions'])
        self.targets.append(batch_targets)

    def get_stack(self):
        predictions = torch.vstack(self.predictions)
        targets = torch.vstack(self.targets)
        predictions = {'predictions': predictions}
        return predictions, targets

class Saver:
    def __init__(self, metric_name, is_lower_better, config, save_name, should_save=True):
        self.metric_name = metric_name
        self.save_path = Path(f'{save_name}/{metric_name}')
        self.operator_function = operator.le if is_lower_better else operator.ge
        self.best_score = np.inf if is_lower_better else 0
        self.best_iteration_num = 0
        self.config = config
        self.save_config()
        self.should_save = should_save

    def update(self, current_iteration_num, current_score, model, tokenizer):
        if self.operator_function(current_score, self.best_score):
            self.best_score = current_score
            self.best_iteration_num = current_iteration_num
            if self.should_save:
                self.save(model, tokenizer)
            print(f'{self.metric_name} attained best score: {current_score:.3f}. Saving the model')

    def save_config(self):
        shutil.rmtree(self.save_path, ignore_errors=True)
        os.makedirs(self.save_path)
        with open(self.save_path / 'config.json', 'w') as fp:
            json.dump(self.config, fp, sort_keys=True, indent=4)

    def save(self, model, tokenizer):
        shutil.rmtree(self.save_path, ignore_errors=True)
        os.makedirs(self.save_path)
        torch.save(model.state_dict(), self.save_path / 'model.pth')
        tokenizer.save_pretrained(self.save_path)

    def get_best_score(self):
        return {'best_score': self.best_score, 'best_iteration_num': self.best_iteration_num}

def train_one_batch(iteration_num, batch, tokenizer, model, optimizer, scheduler, forward_pass_fn, compute_loss_fn, max_length, accumulation_steps, device, **kwargs):
    model.train()
    if iteration_num == 0:
        optimizer.zero_grad()
    batch_loss = forward_pass_fn(batch=batch, tokenizer=tokenizer, model=model,
                                 compute_loss_fn=compute_loss_fn, max_length=max_length,
                                 device=device, **kwargs)['loss']                                  # Forward pass
    batch_loss = batch_loss / accumulation_steps                                                   # Normalize our loss (if averaged)
    batch_loss.backward()                                                                          # Backward pass
    if (iteration_num + 1) % accumulation_steps == 0:                                              # Wait for several backward steps
        optimizer.step()                                                                           # Now we can do an optimizer step
        optimizer.zero_grad()
        if scheduler is not None:
            scheduler.step()
    return model, batch_loss * accumulation_steps

def evaluate(dataloader, tokenizer, model, forward_pass_fn, compute_loss_fn, compute_metric_fn, stacker_class, max_length, device, **kwargs):
    epoch_loss = 0
    model.eval()
    stacker = stacker_class()
    with torch.no_grad():
        for batch_num, batch in enumerate(dataloader):
            batch_outputs = forward_pass_fn(batch=batch, tokenizer=tokenizer, model=model,
                                            compute_loss_fn=compute_loss_fn, max_length=max_length,
                                            device=device, **kwargs)
            batch_loss = batch_outputs['loss']
            epoch_loss += batch_loss.item()
            stacker.update(batch_outputs=batch_outputs, batch_targets=batch['target'])
    average_epoch_loss = epoch_loss/(batch_num+1)
    outputs, targets = stacker.get_stack()
    metric_score = compute_metric_fn(outputs=outputs, targets=targets, **kwargs)
    return average_epoch_loss, metric_score

def train_and_evaluate(num_epochs, train_dataloader, valid_dataloader, tokenizer, model, optimizer, scheduler,
                       forward_pass_fn_train, forward_pass_fn_valid, compute_loss_fn_train, compute_loss_fn_valid,
                       compute_metric_fn, stacker_class, max_length, accumulation_steps, validate_every_n_iteraion,
                       valid_loss_saver, valid_score_saver, device, **kwargs):
    iteration_num = 0
    for epoch_num in range(num_epochs):
        for batch in train_dataloader:
            # for param_group in optimizer.param_groups:
            #     wandb.log({param_group['name']: {"lr": param_group['lr']}})
            model, iteration_train_loss = train_one_batch(iteration_num=iteration_num, batch=batch, tokenizer=tokenizer, model=model,
                                                          optimizer=optimizer, scheduler=scheduler, forward_pass_fn=forward_pass_fn_train,
                                                          compute_loss_fn=compute_loss_fn_train, max_length=max_length,
                                                          accumulation_steps=accumulation_steps, device=device, **kwargs)
            print(f'Epoch_num: {epoch_num}, iteration_num: {iteration_num}, iteration_train_loss: {iteration_train_loss}')
            # wandb.log({'Epoch_num': epoch_num, 'iteration_num': iteration_num, 'iteration_train_loss': iteration_train_loss})

            if 'validate_after_n_iteration' in kwargs:
                validate_after_n_iteration = kwargs['validate_after_n_iteration']
            else:
                validate_after_n_iteration = -1

            if ((iteration_num + 1) % validate_every_n_iteraion == 0) and (iteration_num > validate_after_n_iteration):
                valid_loss, valid_score = evaluate(dataloader=valid_dataloader, tokenizer=tokenizer, model=model,
                                                   forward_pass_fn=forward_pass_fn_valid, compute_loss_fn=compute_loss_fn_valid,
                                                   compute_metric_fn=compute_metric_fn, stacker_class=stacker_class,
                                                   max_length=max_length, device=device, **kwargs)
                valid_loss_saver.update(current_iteration_num=iteration_num,current_score=valid_loss, model=model, tokenizer=tokenizer)
                valid_score_saver.update(current_iteration_num=iteration_num, current_score=valid_score, model=model, tokenizer=tokenizer)
                # wandb.log({"iteration_num": iteration_num, "valid_loss": valid_loss, "valid_score": valid_score})
                print(f'Epoch_num: {epoch_num}, iteration_num: {iteration_num}, iteration_train_loss: {iteration_train_loss}')
                print(f'Epoch_num: {epoch_num}, iteration_num: {iteration_num}, valid_loss: {valid_loss}, valid_score: {valid_score}')
                # wandb.run.summary["best_valid_loss_iteration_num"] = valid_loss_saver.get_best_score()['best_iteration_num']
                # wandb.run.summary["best_valid_loss"] = valid_loss_saver.get_best_score()['best_score']
                # wandb.run.summary["best_valid_score_iteration_num"] = valid_score_saver.get_best_score()['best_iteration_num']
                # wandb.run.summary["best_valid_score"] = valid_score_saver.get_best_score()['best_score']
            iteration_num += 1

    if 'final_model_saver' in kwargs:
        final_model_saver = kwargs['final_model_saver']
        valid_score = 1
        final_model_saver.update(current_iteration_num=iteration_num, current_score=valid_score, model=model, tokenizer=tokenizer)

    # wandb.run.summary["best_valid_loss_iteration_num"] = valid_loss_saver.get_best_score()['best_iteration_num']
    # wandb.run.summary["best_valid_loss"] = valid_loss_saver.get_best_score()['best_score']
    # wandb.run.summary["best_valid_score_iteration_num"] = valid_score_saver.get_best_score()['best_iteration_num']
    # wandb.run.summary["best_valid_score"] = valid_score_saver.get_best_score()['best_score']

    output = {'model': model, 'best_score': valid_score_saver.get_best_score()['best_score'], 'best_loss': valid_loss_saver.get_best_score()['best_score']}
    return output


# Digit fifth powers
"""
Surprisingly there are only three numbers that can be written as the sum of fourth
powers of their digits:

1634 = 1**4 + 6**4 + 3**4 + 4**4
8208 = 8**4 + 2**4 + 0**4 + 8**4
9474 = 9**4 + 4**4 + 7**4 + 4**4
As 1 = 1**4 is not a sum it is not included.

The sum of these numbers is 1634 + 8208 + 9474 = 19316.

Find the sum of all the numbers that can be written as the sum of fifth powers of their digits.
"""

ans = []

for i in range(10, 1000000):
    digit_sum = 0
    for digit in str(i):
        digit_sum += int(digit) ** 5

    if digit_sum == i:
        ans.append(i)

print(ans)
print(sum(ans))


import torch
import torch.nn as nn

class BaselineArch(nn.Module):
    def __init__(self, in_channels: int, channels: int, stride: int):
        super().__init__()
        self.out_channels = out_channels = channels
        self.relu = nn.ReLU(inplace=True)
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1,
                               stride=stride, bias=False)
        self.norm1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1,
                               stride=1, bias=False)
        self.norm2 = nn.BatchNorm2d(out_channels)

        if in_channels == out_channels and stride == 1:
            self.shortcut = nn.Sequential()
        else:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, 1, padding=0,
                          stride=stride, bias=False),
                nn.BatchNorm2d(out_channels)
            )

    def forward(self, x):
        y = self.conv1(x)
        y = self.norm1(y)
        y = self.relu(y)
        y = self.conv2(y)
        y = self.norm2(y)
        y += self.shortcut(x)
        y = self.relu(y)
        return y

class BottleneckArch(nn.Module):
    def __init__(self, in_channels: int, channels: int, stride: int):
        super().__init__()
        self.out_channels = out_channels = channels * 4
        self.relu = nn.ReLU(inplace=True)
        self.conv1 = nn.Conv2d(in_channels, channels, 1, padding=0,
                               stride=stride, bias=False)
        self.norm1 = nn.BatchNorm2d(channels)
        self.conv2 = nn.Conv2d(channels, channels, 3, padding=1,
                               stride=1, bias=False)
        self.norm2 = nn.BatchNorm2d(channels)
        self.conv3 = nn.Conv2d(channels, out_channels, 1, padding=0,
                               stride=1, bias=False)
        self.norm3 = nn.BatchNorm2d(out_channels)

        if in_channels == out_channels and stride == 1:
            self.shortcut = nn.Sequential()
        else:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, 1, padding=0,
                          stride=stride, bias=False),
                nn.BatchNorm2d(out_channels)
            )

    def forward(self, x):
        y = self.conv1(x)
        y = self.norm1(y)
        y = self.relu(y)
        y = self.conv2(y)
        y = self.norm2(y)
        y = self.relu(y)
        y = self.conv3(y)
        y = self.norm3(y)
        y += self.shortcut(x)
        y = self.relu(y)
        return y

class Layer(nn.Module):
    def __init__(self, Block, in_channels: int, channels: int, stride: int, count: int):
        super().__init__()
        blocks = []

        block = Block(in_channels, channels, stride)
        blocks.append(block)
        in_channels = block.out_channels

        for i in range(1, count):
            block = Block(in_channels, channels, 1)
            blocks.append(block)
            in_channels = block.out_channels

        self.blocks = nn.Sequential(*blocks)
        self.out_channels = in_channels

    def forward(self, x):
        return self.blocks(x)

class ResNet(nn.Module):
    def __init__(self, Block, counts: int, start_channels: int, classes: int):
        super().__init__()
        self.relu = nn.ReLU(inplace=True)
        self.conv0 = nn.Conv2d(3, start_channels, 7, padding=3,
                               stride=2, bias=False)
        self.norm0 = nn.BatchNorm2d(start_channels)
        self.pool0 = nn.MaxPool2d(3, stride=2, padding=1)
        self.layer1 = Layer(Block, start_channels,
                            start_channels, 1, counts[0])
        self.layer2 = Layer(Block, self.layer1.out_channels,
                            start_channels*2, 2, counts[1])
        self.layer3 = Layer(Block, self.layer2.out_channels,
                            start_channels*4, 2, counts[2])
        self.layer4 = Layer(Block, self.layer3.out_channels,
                            start_channels*8, 2, counts[3])
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.linear = nn.Linear(self.layer4.out_channels, classes)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode="fan_out",
                                        nonlinearity="relu")
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def forward(self, x):
        x = self.conv0(x)
        x = self.norm0(x)
        x = self.relu(x)
        x = self.pool0(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.linear(x)
        return x

def resnet18(start_channels: int, classes: int):
    return ResNet(BaselineArch, [2, 2, 2, 2], start_channels, classes)

def resnet34(start_channels: int, classes: int):
    return ResNet(BaselineArch, [3, 4, 6, 3], start_channels, classes)

def resnet50(start_channels: int, classes: int):
    return ResNet(BottleneckArch, [3, 4, 6, 3], start_channels, classes)

def resnet101(start_channels: int, classes: int):
    return ResNet(BottleneckArch, [3, 4, 23, 3], start_channels, classes)

def resnet152(start_channels: int, classes: int):
    return ResNet(BottleneckArch, [3, 8, 36, 3], start_channels, classes)


class Solution:
    def permute(self, nums: List[int]) -> List[List[int]]:
        self.res = []
        self.nums = nums
        for n in self.nums:
            self.backtrace([],n)
        return self.res

    def backtrace(self, temp, n):
        temp.append(n)
        if len(temp) == len(self.nums):
            self.res.append(temp)
        else:
            for n in self.nums:
                if n not in temp:
                    self.backtrace(temp.copy(),n)


str1 = '''
1. add
2. sub
3. quit
'''
while True:
    print(str1)
    number = int(input('번호를 입력하세요 : '))
    if number == 1:
        print('\n덧셈')
        num1 = int(input('첫번째 숫자 : '))
        num2 = int(input('두번째 숫자 : '))
        result = num1 + num2
        print(f'{num1}과 {num2}의 합은 {result}입니다.')
        break
    elif number == 2:
        print('뺄셈')
        num1 = int(input('첫번째 숫자 : '))
        num2 = int(input('두번째 숫자 : '))
        result = num1 - num2
        print(f'{num1}에서 {num2}을 빼면 {result}입니다. ')
        break
    elif number == 3:
        print('종료')
        break
    else:
        print('잘못 입력됬습니다.')


"""
주소: https://leetcode.com/problems/count-servers-that-communicate/

내용
- 서버들이 2차원 grid 형태의 망에 붙어있고, 한 서버는 자기와 직선으로 연결된 서버만 통신할 수 있다
- 서버끼리 중계가 가능할때 다른 서버와 통신이 가능한 서버의 수를 구하라

예제
Example 1:
Input: grid = [[1,0],[0,1]]
Output: 0
Explanation: No servers can communicate with others.

Example 2:
Input: grid = [[1,0],[1,1]]
Output: 3
Explanation: All three servers can communicate with at least one other server.

Example 3:
Input: grid = [[1,1,0,0],[0,0,1,0],[0,0,1,0],[0,0,0,1]]
Output: 4
Explanation: The two servers in the first row can communicate with each other. The two servers in the third column can communicate with each other. The server at right bottom corner can't communicate with any other server.

풀이방법
- 단순 그래프에서 크기가 2 이상인 집합을 찾는 문제다
- 그래프로 만들어서 풀면 더 빠르겠지만 일단 메모리 덜쓰게 (귀찮아서)
- 한 집합을 포함하는 row와 col의 set으로 표현하고
- 집합끼리 전부 비교하며 겹쳐지는건 합친다
- 이후 집합 크기가 2 이상인거들의 크기를 합쳐서 리턴
"""

class Solution:
    def countServers(self, grid: List[List[int]]) -> int:
        def compareTwoGrid(grid1, grid2):
            for r in grid1[0]:
                if r in grid2[0]:
                    return True

            for c in grid1[1]:
                if c in grid2[1]:
                    return True

            return False

        def reduceGrid(grids):
            if len(grids)<=1:
                return grids

            retry = True
            while retry:
                retry = False
                for i in range(0, len(grids)-1):
                    for j in range(i+1, len(grids)):
                        if compareTwoGrid(grids[i], grids[j]):
                            grids[i][0] = grids[i][0] | grids[j][0]
                            grids[i][1] = grids[i][1] | grids[j][1]
                            grids[i][2] += grids[j][2]
                            del grids[j]
                            retry = True
                            break
                    if retry:
                        break

            return grids

        grids = []

        for i in range(len(grid)):
            for j in range(len(grid[i])):
                if grid[i][j]==1:
                    grids.append([set([i]), set([j]), 1])

        reduced = reduceGrid(grids)

        res = 0
        for each in reduced:
            print(each)
            if each[2]>1:
                res += each[2]

        return res


# # 1
# li1 = [1, 2, 3]
# li2 = [4, 5, 6]
# for i in li2:
#     li1.append(i)
# print (li1)

# # 2
# for i in range(1499, 2501):
#     if i%5==0 or i%7==0:
#         print(i)

# # 3
# names = ['Samus', 'Cortana', 'V', 'Link', 'Mario', 'Cortana', 'Samus']

# user_name = input("Give me a name: ").lower().capitalize()
# for i, name in enumerate(names):
#     if name == user_name:
#         print (i)
#         break

# # Exercise 4: Greatest Number
# list_num = list()
# k = 0
# while k < 3:
#     list_num.append(int(input("Give me a number: ")))
#     k+=1
# print (sorted(list_num)[-1])

# # Exercise 5: The Alphabet
# alfabet = "abcdefghijklmnopqrstuvwxwz"
# for char in alfabet:
#     if char in ['a', 'e', 'i', 'o', 'u']:
#         print(f"Character {char} is vowel")
#     else:
#         print(f"Character {char} is consonant")

# # Exercise 6: Words And Letters
# words = list()
# for i in range (7):
#     words.append(input("Give me a word: ").lower())

# letter = input("Give me one character: ").lower()

# for word in words:
#     if letter not in word:
#         print (f"The letter {letter.upper()} doesn't exist in the word '{word}'.")
#     else:
#         for i, char in enumerate(word):
#             if char == letter:
#                 print (f"The first index of letter {letter.upper()} in word '{word}' is {i}.")
#                 break

# # Exercise 7
# list_million = []
# for i in range (1, 1000001):
#     list_million.append(i)
# print(min(list_million))
# print(max(list_million))
# print(sum(list_million))

# # Exercise 8 : List And Tuple
# numbers = '34,67,55,33,12,98'
# list_num = (numbers.split(','))
# print(list_num)
# tuple_num = tuple(numbers.split(','))
# print(tuple_num)

# Exercise 9 : Random Number
import random
loses = 0
wins = 0
while True:
    user_num = input("")
    compare_num = random.randint(1,9)
    if user_num == "E":
        print(f"Loses: {loses}, wins: {wins}")
        break
    elif user_num == compare_num:
        print ("Winner!")
        wins += 1
    else:
        print ("Better luck next time")
        loses+=1


import socket
import threading

addres = ('localhost', 12345)

def handle_client(connection,addres):
    print(f'Nova conexao de: {addres}')

    while True:
        data = connection.recv(1024)
        with open("prova_recebido.txt", 'wb') as arquivo:
            arquivo.write(data)
            arquivo.close()

        arquivo1 = open("prova_recebido.txt", "r")
        arquivo2 = open("gabarito.txt", "r")

        f1 = arquivo1.readlines()
        f2 = arquivo2.readlines()

        correto = 0
        errado = 0

        arquivo1.close()
        arquivo2.close()

        for i in range(0, len(f2)):
            for j in range(0, len(f1)):
                if (f2[i][0] == f1[j][0]):
                    if (f1[j] == f2[i]):
                        correto += 1
                    else:
                        errado += 1

        if(i > j-1):
            errado += i-(j-1)

        connection.sendall(str.encode(f'Nome: {f1[0][:-1]} | Acertos: {correto} | Erros: {errado}'))

server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
server.bind((addres))
server.listen()
print('Aguardando conexao...')

while True:
    connection, address = server.accept()
    thread = threading.Thread(target=handle_client, args=(connection, address))
    thread.start()
    print(f"CONEXOES ATIVAS: {threading.activeCount() - 1}")


class vertex:
    def __init__(self,n):
        self.name=n
        self.nbr=list()
    def add_nbr(self,v):
        if v not in self.nbr:
            self.nbr.append(v)
            self.nbr.sort()

class graph:
    vertices={}
    def add_v(self,vx):
        if isinstance(vx,vertex) and vx.name not in self.vertices:
            self.vertices[vx.name]=vx
            return True
        else:
            return False
    def add_e(self,u,v):
        if u in self.vertices and v in self.vertices:
            for key,value in self.vertices.items():
                if key==u:
                    value.add_nbr(v)
                if key==v:
                    value.add_nbr(u)
            return True
        else:
            return False
    def print(self):
        for key in sorted(list(self.vertices.keys())):
            print(key+str(self.vertices[key].nbr))

g=graph()
a=vertex('a')
g.add_v(a)
b=vertex('b')
g.add_v(b)

g.add_e('a','b')

g.print()


import pandas as pd
import json
from collections import Counter
import nltk

class BigramFreq(object):
    """
    Calculates the frequency a word occurs within a given subreddit.
    WordFreq() deals with all time word freq, WordFreqByDay() gets top
    25 words used within a given day.
    """

    def __init__(self, comments, posts):
        # disable SettingWithCopyWarning
        # https://stackoverflow.com/a/20627316/4254021
        pd.options.mode.chained_assignment = None

        self.comments = comments
        self.posts = posts
        self.bigram_freq_df = None

    def get_bigram_freq(self, old_bigrams):
        """
        Creating a list of bigram frequencies from the gathered
        reddit comments and posts
        """
        # Copy of dataframes
        comments = self.comments.copy()
        posts = self.posts.copy()

        # Respective counters
        comment_counts = Counter()
        post_counts = Counter()

        # Loop through both dataframes, finding the most frequent bigrams
        for sent in comments["Text"]:
            words = nltk.word_tokenize(sent)
            comment_counts.update(nltk.bigrams(words))
        bc = pd.DataFrame(list(comment_counts.items()), columns=['bigram', 'n_comment'])

        for sent in posts["Text"]:
            words = nltk.word_tokenize(sent)
            post_counts.update(nltk.bigrams(words))
        bp = pd.DataFrame(list(post_counts.items()), columns=['bigram', 'n_post'])

        # Merge both dataframes and join both words to create bigram
        bigram = pd.merge(bc, bp, on='bigram', how='outer')
        bigram.fillna(0, inplace=True)
        bigram['n'] = bigram['n_comment'] + bigram['n_post']
        bigram["bigram"] = bigram["bigram"].apply(lambda x: ' '.join(x))

        self.bigram_freq_df = bigram
        # sort and limit to 500 bigrams
        bigram = bigram.sort_values('n', ascending=False).head(500)
        # If old bigram count in database, merge new with old and return
        # the updated counts
        if old_bigrams != None:
            old_bigrams = pd.DataFrame.from_records(data=old_bigrams)
            oldnew_merged = pd.concat([bigram,old_bigrams], sort=True)
            # updated counts
            bigram = oldnew_merged.groupby('bigram').sum().reset_index()

        bigram = bigram.to_json(orient='records', date_format=None)
        return json.loads(bigram)

    def get_bigram_freq_by_day(self, oldbigrams):
        # copy dataframes
        comments = self.comments.copy()
        posts = self.posts.copy()
        # Change date format from timestamp
        comments["Date"] = pd.to_datetime(comments['Date'], errors='coerce')
        posts["Date"] = pd.to_datetime(posts['Date'], errors='coerce')
        # Convert back to string date without hours
        comments["Date"] = comments["Date"].dt.strftime('%Y-%m-%d')
        posts["Date"] = posts["Date"].dt.strftime('%Y-%m-%d')
        #Merge dataframes
        merged = pd.concat([posts,comments], sort=True)
        # Group by date
        merged = merged.groupby('Date')
        bbd_list = []
        # loop through groups
        for name, group in merged:
            merged_counts = Counter()
            # Loop through text counting bigrams
            for sent in group["Text"]:
                words = nltk.word_tokenize(sent)
                merged_counts.update(nltk.bigrams(words))
            updated = []
            # join bigrams to one word
            for key, value in merged_counts.most_common(25):
                k = ' '.join(key)
                updated.append({
                    "bigram": k,
                    "n": value
                })
            # append grouped date and counted bigrams to list
            bbd_list.append([name, updated])

        # Create dataframe with counts
        bbd = pd.DataFrame(bbd_list, columns = ['Date','counts'])

        if oldbigrams != None:
            oldbigrams = pd.DataFrame.from_records(data=oldbigrams)
            oldnew_merged = pd.concat([bbd,oldbigrams], sort=True)

            flattened = []
            # loop through groups
            for date, group in oldnew_merged.groupby('Date'):
                for array in group["counts"]:
                    for obj in array:
                        flattened.append([date, obj["bigram"], obj["n"]])

            # Goal is to group similar dates and marge words together
            # keeping the 25 most popular for each date
            flattened = pd.DataFrame(flattened, columns = ['Date','bigram','n'])
            flattened = flattened.groupby(['Date', 'bigram'])['n'].sum().reset_index()

            flat = []
            for date, group in flattened.groupby('Date'):
                bigram = list(group["bigram"])
                n = list(group["n"])
                group_arr = []
                for i in range(len(n)-1):
                    if i > 25:
                        break
                    group_arr.append({
                        "bigram":bigram[i],
                        "n":n[i]
                    })
                flat.append([date, group_arr])
                group_arr = []

            bbd = pd.DataFrame(flat, columns = ['Date','counts'])

        bbd = bbd.to_json(orient='records', date_format=None)
        bbd = json.loads(bbd)
        return bbd


from database import Database

class CheckingAccount:
    # Initializes a checking account
    def __init__(self, database, rg, client_name):
        self._database = database
        self.rg = rg
        self.client_name = client_name

        self.__create(rg, client_name)

    # Persists a new checking account
    def __create(self, rg, client_name):
        checking_account = self._database.get_checking_account(rg)

        if (checking_account is not None):
            return

        self._database.create(rg, client_name)

    # Verifies if checking account exists
    @staticmethod
    def login(database, rg):
        checking_account = database.get_checking_account(rg)

        if (checking_account is None):
            return None

        rg, client_name, _ = checking_account

        return CheckingAccount(database, rg, client_name)

    # Get balance
    def balance(self):
        _, _, amount_in_cents = self._database.get_checking_account(self.rg)

        return amount_in_cents

    # Withdraws from current checking account
    def withdraw(self, amount_in_cents):
        _, _, ck_amount_in_cents = self._database.get_checking_account(
            self.rg)

        if (amount_in_cents <= 0):
            raise Exception("Invalid amount to be withdrawn")

        if (amount_in_cents > ck_amount_in_cents):
            raise Exception(
                "Not enough funds to be withdrawn: {}".format(ck_amount_in_cents / 100))

        self._database.withdraw_atomic(self.rg, amount_in_cents)

    # Deposits to current checking account
    def deposit(self, amount_in_cents):
        if (amount_in_cents <= 0):
            raise Exception("Invalid amount to be deposited")

        self._database.deposit_atomic(self.rg, amount_in_cents)

    # Transfers to another checking account
    def transfer_to(self, destination_rg, amount_in_cents):
        if (amount_in_cents <= 0):
            raise Exception("Invalid amount to be transfered")

        destination_ck = self._database.get_checking_account(destination_rg)

        if (destination_ck is None):
            raise Exception("Unable to make transfer, destination not exists")

        _, _, origin_amount_in_cents = self._database.get_checking_account(
            self.rg)

        if (amount_in_cents > origin_amount_in_cents):
            raise Exception("Not enough funds: {}".format(
                origin_amount_in_cents / 100))

        self._database.withdraw(self.rg, amount_in_cents)
        self._database.deposit(destination_rg, amount_in_cents)
        self._database.save()

# rg1 = 12312312
# rg2 = 12312313
#
# db = Database.instance()
# ck1 = CheckingAccount(db, rg1, "Teste 1")
# ck2 = CheckingAccount(db, rg2, "Teste 2")
#
# ck1.deposit(10)
# ck1.withdraw(5)
# ck1.transfer_to(rg2, 3)


from inspect import isfunction

import matplotlib.pyplot as plt
import numpy as np
from einops import rearrange, reduce
from torchvision import transforms
from torchvision.transforms import (
    CenterCrop,
    Compose,
    Lambda,
    Resize,
    ToPILImage,
    ToTensor,
)

import os
import random

import torch
import json

def exists(x):
    return x is not None

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

def num_to_groups(num, divisor):
    groups = num // divisor
    remainder = num % divisor
    arr = [divisor] * groups
    if remainder > 0:
        arr.append(remainder)
    return arr

def plot(imgs, with_orig=False, row_title=None, **imshow_kwargs):
    if not isinstance(imgs[0], list):
        # Make a 2d grid even if there's just 1 row
        imgs = [imgs]

    num_rows = len(imgs)
    num_cols = len(imgs[0]) + with_orig
    fig, axs = plt.subplots(
        figsize=(200, 200), nrows=num_rows, ncols=num_cols, squeeze=False
    )
    for row_idx, row in enumerate(imgs):
        # row = [image] + row if with_orig else row
        for col_idx, img in enumerate(row):
            ax = axs[row_idx, col_idx]
            ax.imshow(np.asarray(img), **imshow_kwargs)
            ax.set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])

    if with_orig:
        axs[0, 0].set(title="Original image")
        axs[0, 0].title.set_size(8)
    if row_title is not None:
        for row_idx in range(num_rows):
            axs[row_idx, 0].set(ylabel=row_title[row_idx])

    plt.tight_layout()

def num_to_groups(num, divisor):
    groups = num // divisor
    remainder = num % divisor
    arr = [divisor] * groups
    if remainder > 0:
        arr.append(remainder)
    return arr

def extract(a, t, x_shape):
    batch_size = t.shape[0]
    out = a.gather(-1, t.cpu())
    return out.reshape(batch_size, *((1,) * (len(x_shape) - 1))).to(t.device)

def get_transform_aux(image_size):
    transform = Compose(
        [
            Resize(image_size),
            CenterCrop(image_size),
            ToTensor(),  # turn into torch Tensor of shape CHW, divide by 255
            Lambda(lambda t: (t * 2) - 1),
        ]
    )
    return transform

transform = Compose(
    [
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Lambda(lambda t: (t * 2) - 1),
    ]
)

reverse_transform = Compose(
    [
        Lambda(lambda t: (t + 1) / 2),
        Lambda(lambda t: t.permute(1, 2, 0)),  # CHW to HWC
        Lambda(lambda t: t * 255.0),
        Lambda(lambda t: t.numpy().astype(np.uint8)),
        ToPILImage(),
    ]
)

class Settings:
    def __init__(self, **kwargs):
        for k, v in kwargs.items():
            self.__dict__[k] = v

    def __getitem__(self, key):
        return self.__dict__[key]

    def __setitem__(self, key, val):
        self.__setattr__(key, val)

    def __repr__(self):
        return json.dumps(self.__dict__, indent=4, default=lambda x: x.__repr__())

def plot_images(samples, settings: Settings, n=64):
    sqrt = int(n**0.5)
    fig, ax = plt.subplots(8, 8, figsize=(20, 20))
    ax = ax.ravel()
    for i in range(n):
        ax[i].imshow(
            (samples[-1][i].reshape(
                settings.image_size, settings.image_size, settings.channels
            ) + 1) / 2,
            # cmap="gray",
        )
    plt.tight_layout()
    plt.show()

def set_all_seeds(seed=0):
    # python's seeds
    random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)
    np.random.seed(seed)

    # torch's seeds
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


import os

import numpy as np
import scipy.io as sio
from PIL import Image

from data.utils import bbox
from data.utils import transform
from data.utils.camera import Camera
from data.utils.data_utils import TEST_MODE, TRAIN_MODE, VALID_MODE, get_visible_joints

from model.utils import pose_format

def load_h36m_mat_annotation(filename):
    mat = sio.loadmat(filename, struct_as_record=False, squeeze_me=True)

    # Respect the order of TEST (0), TRAIN (1), and VALID (2)
    sequences = [mat['sequences_te'], mat['sequences_tr'], mat['sequences_val']]
    action_labels = mat['action_labels']
    joint_labels = mat['joint_labels']

    return sequences, action_labels, joint_labels

def serialize_index_sequences(seq):
    frames_idx = []
    for s in range(len(seq)):
        for f in range(len(seq[s].frames)):
            frames_idx.append((s, f))

    return frames_idx

def get_clip_frame_index(
    sequence_size,
    subsample,
    num_frames,
    random_clip=False):

    # Assert that subsample is integer and positive
    assert (type(subsample) == int) and subsample > 0

    idx_coef = 1.
    while idx_coef * sequence_size < num_frames:
        idx_coef *= 1.5
    sequence_size *= idx_coef

    # Check if the given subsample value is feasible, otherwise, reduce
    # it to the maximum acceptable value.
    max_subsample = int(sequence_size / num_frames)
    if subsample > max_subsample:
        subsample = max_subsample

    vidminf = subsample * (num_frames - 1) + 1 # Video min num of frames
    maxs = sequence_size - vidminf # Maximum start
    if random_clip:
        start = np.random.randint(maxs + 1)
    else:
        start = int(maxs / 2)

    frames = list(range(start, start + vidminf, subsample))
    if idx_coef > 1:
        for i in range(len(frames)):
            frames[i] = int(frames[i] / idx_coef)

    return frames

class Human36M(object):
    '''
    Human3.6M dataset for 3D pose estimation dataset manager
    '''

    def __init__(self,
                 dataset_path,
                 dataconf,
                 poselayout=pose_format.pa17j3d,
                 topology='sequences',
                 clip_size=16):

        assert topology in ['sequences', 'frames'], 'Invalid topology ({})'.format(topology)

        self.dataset_path = dataset_path
        self.dataconf = dataconf
        self.poselayout = poselayout
        self.topology = topology
        self.clip_size = clip_size
        self.load_annotations(os.path.join(dataset_path, 'annotations.mat'))

    def load_annotations(self, filename):
        try:
            # sequences here is a list of 3: test sequences, train sequences, val sequences
            self.sequences, self.action_labels, self.joint_labels = load_h36m_mat_annotation(filename)
            self.frame_idx = [serialize_index_sequences(self.sequences[0]),
                    serialize_index_sequences(self.sequences[1]),
                    serialize_index_sequences(self.sequences[2])]

            global ACTION_LABELS
            ACTION_LABELS = self.action_labels

        except:
            warning('Error loading Human3.6M dataset!')
            raise

    def get_data(self, key, mode, frame_list=None, fast_crop=False):
        '''
        pose_w:     ground truth data in millimeters
                    loaded from the annotation file and formatted to the self.poselayout format (= picking out joints)
        pose_uvd:   same but with xy coords given in pixels (uv) and depth (d) in millimeters
        pose:       same format as pose_uvd but with all transformations (crop, hflip, etc) applied
        '''
        output = {}

        if mode == TRAIN_MODE:
            dconf = self.dataconf.random_data_generator()
            random_clip = True
        else:
            dconf = self.dataconf.get_fixed_config()
            random_clip = False

        if self.topology == 'sequences':
            seq = self.sequences[mode][key]
            if frame_list is None:
                frame_list = get_clip_frame_index(len(seq.frames),
                        dconf['subspl'], self.clip_size,
                        random_clip=random_clip)
            objframes = seq.frames[frame_list]
        else:
            seq_idx, frame_idx = self.frame_idx[mode][key]
            seq = self.sequences[mode][seq_idx]
            objframes = seq.frames[[frame_idx]]

        # Build a Camera object
        cpar = seq.camera_parameters
        cam = Camera(cpar.R, cpar.T, cpar.f, cpar.c, cpar.p, cpar.k)

        # Load and project the poses
        pose_w = self.load_pose_annot(objframes)
        pose_uvd = cam.project(np.reshape(pose_w, (-1, 3)))
        pose_uvd = np.reshape(pose_uvd, (len(objframes), self.poselayout.num_joints, 3))

        # Compute GT bouding box
        imgsize = (objframes[0].w, objframes[0].h)
        objpos, winsize, zrange = bbox.get_crop_params(pose_uvd[:, 0, :],
                imgsize, cam.f, dconf['scale'])

        objpos += dconf['scale'] * np.array([dconf['transx'], dconf['transy']])
        frames = np.empty((len(objframes),) + self.dataconf.input_shape)
        pose = np.empty((len(objframes), self.poselayout.num_joints, self.poselayout.dim))

        for i in range(len(objframes)):
            image = 'images/%s/%05d.jpg' % (seq.name, objframes[i].f)
            imgt = transform.ImageTransform(Image.open(os.path.join(self.dataset_path, image)))

            imgt.rotate_crop(dconf['angle'], objpos, winsize)
            if dconf['hflip'] == 1:
                imgt.horizontal_flip()

            imgt.resize(self.dataconf.crop_resolution)
            imgt.normalize_affinemap()
            frames[i, :, :, :] = transform.normalize_channels(imgt.asarray(),
                    channel_power=dconf['chpower'])

            pose[i, :, 0:2] = transform.transform_2d_points(imgt.afmat, pose_uvd[i, :,0:2], transpose=True)
            pose[i, :, 2] = (pose_uvd[i, :, 2] - zrange[0]) / (zrange[1] - zrange[0])

            if imgt.hflip:
                pose[i, :, :] = pose[i, self.poselayout.map_hflip, :]

        # Set outsider body joints to invalid (-1e9)
        pose = np.reshape(pose, (-1, self.poselayout.dim))
        pose[np.isnan(pose)] = -1e9
        v = np.expand_dims(get_visible_joints(pose[:,0:2]), axis=-1)
        pose[(v==0)[:,0],:] = -1e9
        pose = np.reshape(pose, (len(objframes), self.poselayout.num_joints, self.poselayout.dim))
        v = np.reshape(v, (len(objframes), self.poselayout.num_joints, 1))

        pose = np.concatenate((pose, v), axis=-1)
        if self.topology != 'sequences':
            pose_w = np.squeeze(pose_w, axis=0)
            pose_uvd = np.squeeze(pose_uvd, axis=0)
            pose = np.squeeze(pose, axis=0)
            frames = np.squeeze(frames, axis=0)

        output['camera'] = cam.serialize()
        output['action'] = int(seq.name[1:3]) - 1
        output['pose_w'] = pose_w
        output['rootz'] = pose_uvd[0,2]
        output['pose_uvd'] = pose_uvd
        output['pose'] = pose
        output['frame'] = frames

        # Take the last transformation matrix, it should be the same for all frames
        output['afmat'] = imgt.afmat.copy()

        return output

    def load_pose_annot(self, frames):
        p = np.empty((len(frames), self.poselayout.num_joints, self.poselayout.dim))
        for i in range(len(frames)):
            p[i,:] = frames[i].pose3d.T[self.poselayout.map_from_h36m, 0:self.poselayout.dim].copy()

        return p

    def clip_length(self):
        if self.topology == 'sequences':
            return self.clip_size
        else:
            return None

    def clip_shape(self):
        if self.topology == 'sequences':
            return (self.clip_size,)
        else:
            return ()

    def get_shape(self, dictkey):
        if dictkey == 'frame':
            return self.clip_shape() + self.dataconf.input_shape
        if dictkey == 'pose':
            return self.clip_shape() + (self.poselayout.num_joints, self.poselayout.dim+1)
        if dictkey == 'pose_w':
            return self.clip_shape() + (self.poselayout.num_joints, self.poselayout.dim)
        if dictkey == 'pose_uvd':
            return self.clip_shape() + (self.poselayout.num_joints, self.poselayout.dim)
        if dictkey == 'action':
            return (1,)
        if dictkey == 'camera':
            return (21,)
        if dictkey == 'afmat':
            return (3, 3)
        raise Exception('Invalid dictkey %s on get_shape!' % dictkey)

    def get_length(self, mode):
        if self.topology == 'sequences':
            return len(self.sequences[mode])
        else:
            return len(self.frame_idx[mode])


import random

#get user name
print("BOT: What do you want me to call you?")
user_name = input()

#templates
bot_template = "BOT : {0}"
user_template = user_name + " : {0}"

#defining variables
name = "Bad bot"
weather = "rainy"
mood = "grumpy"

#creating responses
responses = {
    "what's your name?" : [
        "You may call me Mr {0}".format(name),
        "I usually go by {0}".format(name)],
    "what's today's weather?": [
        "Today its {0}".format(weather),
        "It's {0} today".format(weather)],
    "how are you?" : [
        "I'm {0} as always".format(mood),
        "{0}! How about you?".format(mood)],
    "how old are you?" : [
        "don't you know its rude to ask that?"],
    "default": [
        "Does not compute"]
}

#defining response method
def respond(message):
    if message in responses:
        bot_message = random.choice(responses[message])
    else:
        bot_message = random.choice(responses["default"])
    return bot_message

#function for related responses
def related(x_text):
    if "name" in x_text:
        y_text = "what's your name?"
    elif "weather" in x_text:
        y_text = "what's today's weather?"
    elif "how are" in x_text:
        y_text = "how are you?"
    else:
        y_text = ""
    return y_text

#sending message function
def send_message(message):
  print(user_template.format(message))
  response = respond(message)
  print(bot_template.format(response))

#testing
while 1:
    my_input = input()
    my_input = my_input.lower()
    related_text = related(my_input)
    send_message(related_text)
    if my_input == "exit" or my_input == "stop":
        break


exam_hour = int(input())
exam_minute = int(input())
arrival_hour = int(input())
arrival_minute = int(input())

exam_time = exam_hour * 60 + exam_minute
arrival_time = arrival_hour * 60 + arrival_minute

if arrival_time > exam_time:
    print("Late")

    difference = arrival_time - exam_time
    if difference < 60:
        print(f"{difference} minutes after the start")
    else:
        hours = difference // 60
        minutes = difference % 60
        print(f"{hours}:{minutes:02d} hours after the start")
elif exam_time - arrival_time <= 30:
    print("On time")
    # Изчисляване на времето преди началото в минути
    difference = exam_time - arrival_time
    if difference > 0:
        print(f"{difference} minutes before the start")
elif exam_time - arrival_time > 30:
    print("Early")

    difference = exam_time - arrival_time
    if difference < 60:
        print(f"{difference} minutes before the start")
    else:
        hours = difference // 60
        minutes = difference % 60
        print(f"{hours}:{minutes:02d} hours before the start")

Conditional Statements Advanced Exercise/P08on_time_for_the_exam.py


from __future__ import annotations

import time
from dataclasses import dataclass

from alpaca.utils import Visitor
from eisen.common.exceptionshandler import ExceptionsHandler
from eisen.validation.modulevisitor import ModuleVisitor
from eisen.validation.typechecker import TypeChecker
from eisen.validation.functionvisitor import FunctionVisitor
from eisen.validation.usagechecker import UsageChecker
from eisen.validation.declarationvisitor import DeclarationVisitor
from eisen.validation.finalizationvisitor import FinalizationVisitor, Finalization2
from eisen.validation.fnconverter import FnConverter
from eisen.validation.initalizer import Initializer
from eisen.validation.nilcheck import NilCheck
from eisen.validation.instancevisitor import InstanceVisitor
from eisen.validation.vectorvisitor import VectorVisitor
from eisen.trace.memoryvisitor import MemoryVisitor
from eisen.state.basestate import BaseState as State

# Notes:
# A module is a collection of structs/functions
# A context, properly, is a block with instances defined
# A context may have a parent as a module (as module level functions are available
# for use as instances)
# They are implemented the same.

# class used to debug
class Printast():
    def run(self, state: State):
        print(state.get_ast())
        return state

class Workflow():
    steps: list[Visitor] = [
        # initialize the .data attribute for all asts with empty NodeData instances
        Initializer,

        # create the module structure of the program.
        #   - the module of a node can be accessed by params.ast_get_mod()
        ModuleVisitor,

        # add proto types for struct/interfaces, which are the representation
        # of struct declaration without definition.
        DeclarationVisitor,

        # finalizes the proto types (which moves from declaration to definition)
        # TODO: for now, interfaces can only be implemented from the same module
        # in which they are defined.
        #
        # we must finalize interfaces first because structs depend on interfaces
        # as they implement interfaces
        FinalizationVisitor,
        Finalization2,

        # adds types for and constructs the functions. this also normalizes the
        # (def ...) and (create ...) asts so they have the same child structure,
        # which allows us to process them identically later in the
        # TypeClassFlowWrangler.
        FunctionVisitor,
        VectorVisitor,

        # this changes (ref ...) to (fn ...) if they refer to global functions
        FnConverter,

        # evaluate the flow of types through the program.
        #   - the type which is flowed through a node can be accessed by
        #     params.get_returned_type()
        TypeChecker,

        InstanceVisitor,

        # this handles restrictions based on let/var/val differences
        UsageChecker,

        NilCheck,
        MemoryVisitor,
        # MemCheck,
        # MoveVisitor
        # Printast,

        # note: def, create, fn, ref, ilet, :: need instances!!
    ]

    @staticmethod
    def _choose_steps(supplied_steps: list[Visitor]=None):
        if supplied_steps:
            return supplied_steps
        return Workflow.steps

    @staticmethod
    def _run_step_with_state(step: Visitor, state: State) -> State:
        return step().run(state)

    @staticmethod
    def execute(state: State, steps: list[Visitor]=None) -> tuple[bool, State]:
        for step in Workflow._choose_steps(steps):
            state = Workflow._run_step_with_state(step, state)
            ExceptionsHandler().apply(state)
            if Workflow.should_stop_execution(state):
                return False, state
        return True, state

    @staticmethod
    def _instrument_step_with_telemetry(step_times: list[int], step: Visitor) -> Visitor:
        class decorated_step():
            def run(self, state: State):
                start = time.perf_counter_ns()
                state = Workflow._run_step_with_state(step, state)
                end = time.perf_counter_ns()
                decorated_step.add_metric(start, end)
                return state

            @staticmethod
            def add_metric(start: int, end: int):
                step_times.append(PerformanceMetric(
                    step_name=step.__name__,
                    step_time=round((end-start)/1000000, 5)))

        return decorated_step

    @staticmethod
    def execute_with_benchmarks(state: State, steps: list[Visitor]=None) -> tuple[bool, State]:
        metrics: list[PerformanceMetric] = []
        steps_with_telemetry = [Workflow._instrument_step_with_telemetry(metrics, step)
            for step in Workflow._choose_steps(steps)]

        result, state = Workflow.execute(state, steps=steps_with_telemetry)
        PerformanceMetric.print_performance_metrics(metrics)
        return result, state

    @staticmethod
    def should_stop_execution(state: State) -> bool:
        return len(state.exceptions) > 0

@dataclass
class PerformanceMetric():
    step_name: str
    step_time: int

    @staticmethod
    def _print_performance_metric_line(metric: PerformanceMetric, block_size: int):
        print(" "*(block_size - len(metric.step_name)),
            metric.step_name, " ",
            metric.step_time)

    @staticmethod
    def print_performance_metrics(metrics: list[PerformanceMetric]):
        longest_name_size = max(len(metric.step_name) for metric in metrics)
        block_size = longest_name_size + 4
        for metric in metrics:
            PerformanceMetric._print_performance_metric_line(metric, block_size)

        print(" "*(block_size-len("Total")),
              "Total", " ",
              round(sum(metric.step_time for metric in metrics), 5))


from rest_framework import serializers

from .models import Location

# класс позволяющий преобразовать объект класса Location в объект JSON
class LocationSerializer(serializers.ModelSerializer):
    class Meta:
        # в качестве класса модели используется класс Location
        model = Location
        # поля: id, широта, долгота, время добавления
        fields = ('id', 'latitude', 'longitude', 'datetime')
        # данные заносятся через POST только в широту и долготу,
        # id и время выставляются автоматически (поэтому readonly)
        read_only_fields = ('id', 'datetime')

    def create(self, validated_data):
        # создание объекта Location
        # данные заносятся через пост только в широту и долготу,
        # id и время выставляются автоматически
        location = Location.objects.create(
            latitude=validated_data['latitude'],
            longitude=validated_data['longitude']
        )
        # сохранение объекта в БД
        location.save()
        # возврат созданного объекта
        return location


"""Given a string, partition the string into as many parts as possible so that each letter appears in at most one part, and return a list of integers representing the size of these parts.

Example:

Input: s = "ababcbacadefegdehijhklij"
Output: [9,7,8]

Explanation:
The partition is "ababcbaca", "defegde", "hijhklij".
This is a partition so that each letter appears in at most one part.

Time: O(n)
Space: O(26) -> O(1)
"""

def partition_labels(s):
    lookup = {c: i for i, c in enumerate(s)}  # store each char's last index

    start = end = 0

    result = []

    for i, c in enumerate(s):
        end = max(end, lookup[c])

        if i == end:  # found a partition
            result.append(i - start + 1)
            start = i + 1

    return result