from collections import defaultdict

import torch
import torch.nn as nn

import numpy as np

class SplitCrossEntropyLoss(nn.Module):
    r'''SplitCrossEntropyLoss calculates an approximate softmax'''
    def __init__(self, hidden_size, splits, verbose=False):
        # We assume splits is [0, split1, split2, N] where N >= |V|
        # For example, a vocab of 1000 words may have splits [0] + [100, 500] + [inf]
        super(SplitCrossEntropyLoss, self).__init__()
        self.hidden_size = hidden_size
        self.splits = [0] + splits + [100 * 1000000]
        self.nsplits = len(self.splits) - 1
        self.stats = defaultdict(list)
        self.verbose = verbose
        # Each of the splits that aren't in the head require a pretend token, we'll call them tombstones
        # The probability given to this tombstone is the probability of selecting an item from the represented split
        if self.nsplits > 1:
            self.tail_vectors = nn.Parameter(torch.zeros(self.nsplits - 1, hidden_size))
            self.tail_bias = nn.Parameter(torch.zeros(self.nsplits - 1))

    def logprob(self, weight, bias, hiddens, splits=None, softmaxed_head_res=None, verbose=False):
        # First we perform the first softmax on the head vocabulary and the tombstones
        if softmaxed_head_res is None:
            start, end = self.splits[0], self.splits[1]
            head_weight = None if end - start == 0 else weight[start:end]
            head_bias = None if end - start == 0 else bias[start:end]
            # We only add the tombstones if we have more than one split
            if self.nsplits > 1:
                head_weight = self.tail_vectors if head_weight is None else torch.cat([head_weight, self.tail_vectors])
                head_bias = self.tail_bias if head_bias is None else torch.cat([head_bias, self.tail_bias])

            # Perform the softmax calculation for the word vectors in the head for all splits
            # We need to guard against empty splits as torch.cat does not like random lists
            head_res = torch.nn.functional.linear(hiddens, head_weight, bias=head_bias)
            softmaxed_head_res = torch.nn.functional.log_softmax(head_res)

        if splits is None:
            splits = list(range(self.nsplits))

        results = []
        running_offset = 0
        for idx in splits:

            # For those targets in the head (idx == 0) we only need to return their loss
            if idx == 0:
                results.append(softmaxed_head_res[:, :-(self.nsplits - 1)])

            # If the target is in one of the splits, the probability is the p(tombstone) * p(word within tombstone)
            else:
                start, end = self.splits[idx], self.splits[idx + 1]
                tail_weight = weight[start:end]
                tail_bias = bias[start:end]

                # Calculate the softmax for the words in the tombstone
                tail_res = torch.nn.functional.linear(hiddens, tail_weight, bias=tail_bias)

                # Then we calculate p(tombstone) * p(word in tombstone)
                # Adding is equivalent to multiplication in log space
                head_entropy = (softmaxed_head_res[:, -idx]).contiguous()
                tail_entropy = torch.nn.functional.log_softmax(tail_res)
                results.append(head_entropy.view(-1, 1) + tail_entropy)

        if len(results) > 1:
            return torch.cat(results, dim=1)
        return results[0]

    def split_on_targets(self, hiddens, targets):
        # Split the targets into those in the head and in the tail
        split_targets = []
        split_hiddens = []

        # Determine to which split each element belongs (for each start split value, add 1 if equal or greater)
        # This method appears slower at least for WT-103 values for approx softmax
        #masks = [(targets >= self.splits[idx]).view(1, -1) for idx in range(1, self.nsplits)]
        #mask = torch.sum(torch.cat(masks, dim=0), dim=0)
        ###
        # This is equally fast for smaller splits as method below but scales linearly
        mask = None
        for idx in range(1, self.nsplits):
            partial_mask = targets >= self.splits[idx]
            mask = mask + partial_mask if mask is not None else partial_mask
        ###
        #masks = torch.stack([targets] * (self.nsplits - 1))
        #mask = torch.sum(masks >= self.split_starts, dim=0)
        for idx in range(self.nsplits):
            # If there are no splits, avoid costly masked select
            if self.nsplits == 1:
                split_targets, split_hiddens = [targets], [hiddens]
                continue
            # If all the words are covered by earlier targets, we have empties so later stages don't freak out
            if sum(len(t) for t in split_targets) == len(targets):
                split_targets.append([])
                split_hiddens.append([])
                continue
            # Are you in our split?
            tmp_mask = mask == idx
            split_targets.append(torch.masked_select(targets, tmp_mask))
            split_hiddens.append(hiddens.masked_select(tmp_mask.unsqueeze(1).expand_as(hiddens)).view(-1, hiddens.size(1)))
        return split_targets, split_hiddens

    def forward(self, weight, bias, hiddens, targets, verbose=False):
        if self.verbose or verbose:
            for idx in sorted(self.stats):
                print('{}: {}'.format(idx, int(np.mean(self.stats[idx]))), end=', ')
            print()

        total_loss = None
        if len(hiddens.size()) > 2: hiddens = hiddens.view(-1, hiddens.size(2))

        split_targets, split_hiddens = self.split_on_targets(hiddens, targets)

        # First we perform the first softmax on the head vocabulary and the tombstones
        start, end = self.splits[0], self.splits[1]
        head_weight = None if end - start == 0 else weight[start:end]
        head_bias = None if end - start == 0 else bias[start:end]

        # We only add the tombstones if we have more than one split
        if self.nsplits > 1:
            head_weight = self.tail_vectors if head_weight is None else torch.cat([head_weight, self.tail_vectors])
            head_bias = self.tail_bias if head_bias is None else torch.cat([head_bias, self.tail_bias])

        # Perform the softmax calculation for the word vectors in the head for all splits
        # We need to guard against empty splits as torch.cat does not like random lists
        combo = torch.cat([split_hiddens[i] for i in range(self.nsplits) if len(split_hiddens[i])])
        ###
        all_head_res = torch.nn.functional.linear(combo, head_weight, bias=head_bias)
        softmaxed_all_head_res = torch.nn.functional.log_softmax(all_head_res)
        if self.verbose or verbose:
            self.stats[0].append(combo.size()[0] * head_weight.size()[0])

        running_offset = 0
        for idx in range(self.nsplits):
            # If there are no targets for this split, continue
            if len(split_targets[idx]) == 0: continue

            # For those targets in the head (idx == 0) we only need to return their loss
            if idx == 0:
                softmaxed_head_res = softmaxed_all_head_res[running_offset:running_offset + len(split_hiddens[idx])]
                entropy = -torch.gather(softmaxed_head_res, dim=1, index=split_targets[idx].view(-1, 1))
            # If the target is in one of the splits, the probability is the p(tombstone) * p(word within tombstone)
            else:
                softmaxed_head_res = softmaxed_all_head_res[running_offset:running_offset + len(split_hiddens[idx])]

                if self.verbose or verbose:
                    start, end = self.splits[idx], self.splits[idx + 1]
                    tail_weight = weight[start:end]
                    self.stats[idx].append(split_hiddens[idx].size()[0] * tail_weight.size()[0])

                # Calculate the softmax for the words in the tombstone
                tail_res = self.logprob(weight, bias, split_hiddens[idx], splits=[idx], softmaxed_head_res=softmaxed_head_res)

                # Then we calculate p(tombstone) * p(word in tombstone)
                # Adding is equivalent to multiplication in log space
                head_entropy = softmaxed_head_res[:, -idx]
                # All indices are shifted - if the first split handles [0,...,499] then the 500th in the second split will be 0 indexed
                indices = (split_targets[idx] - self.splits[idx]).view(-1, 1)
                # Warning: if you don't squeeze, you get an N x 1 return, which acts oddly with broadcasting
                tail_entropy = torch.gather(torch.nn.functional.log_softmax(tail_res), dim=1, index=indices).squeeze()
                entropy = -(head_entropy + tail_entropy)
            ###
            running_offset += len(split_hiddens[idx])
            total_loss = entropy.float().sum() if total_loss is None else total_loss + entropy.float().sum()

        return (total_loss / len(targets)).type_as(weight)

if __name__ == '__main__':
    np.random.seed(42)
    torch.manual_seed(42)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(42)

    V = 8
    H = 10
    N = 100
    E = 10

    embed = torch.nn.Embedding(V, H)
    crit = SplitCrossEntropyLoss(hidden_size=H, splits=[V // 2])
    bias = torch.nn.Parameter(torch.ones(V))
    optimizer = torch.optim.SGD(list(embed.parameters()) + list(crit.parameters()), lr=1)

    for _ in range(E):
        prev = (torch.rand(N, 1) * 0.999 * V).int().long()
        x = (torch.rand(N, 1) * 0.999 * V).int().long()
        y = embed(prev).squeeze()
        c = crit(embed.weight, bias, y, x.view(N))
        print('Crit', c.exp().item())

        logprobs = crit.logprob(embed.weight, bias, y[:2]).exp()
        print(logprobs)
        print(logprobs.sum(dim=1))

        optimizer.zero_grad()
        c.backward()
        optimizer.step()


"""Creates the network architecture for linearizing PDEs."""
import numpy as np
import tensorflow as tf
from tensorflow import keras

from DenseResBlock import DenseResBlock
from RelMSE import RelMSE

class NetworkArch(keras.Model):
    """Subclass the Keras Model class."""

    def __init__(self,
                 n_inputs=128,
                 n_latent=21,
                 len_time=51,
                 num_shifts=50,
                 num_shifts_middle=50,
                 outer_encoder=DenseResBlock(),
                 outer_decoder=DenseResBlock(),
                 inner_config=dict(),
                 inner_loss_weights=[1, 1],
                 L_diag=False,
                 train_autoencoder_only=False,
                 **kwargs):
        """
        Create network architecture for linearizing PDEs.

        Arguments:
            n_inputs -- the number of inputs to the network
                (spatial discretization of the PDE)
            n_latent -- the dimensionality of the latent space
                (i.e. number of Koopman eigenfunctions in expansion)
            len_time -- number of time steps for each trajectory in data
            num_shifts -- the number of time steps in the future each network
                input will predict when calculating prediction loss
            num_shifts_middle -- the number of time steps in the future each
                network input will predict when calculating the linearity loss
            outer_encoder -- a Keras Layer or Model with the architecture for
                the outer encoder
            outer_decoder -- a Keras Layer or Model with the architecture for
                the outer decoder (typically the same as the outer encoder)
            inner_config -- Python dictionary with keyword arguments for
                the inner encoder and decoder
            inner_loss_weights -- Python list of length 2 with weights for the
                inner autoencoder loss and linearity loss
            L_diag -- Boolean that determines whether the dynamics matrix L is
                constrained to be diagonal
            train_autoencoder_only -- Boolean that determines whether only the
                autoencoder losses are used for training. It is recommendeded
                that you do several epochs of pretraining with only the
                autoencoder losses and then set this option to False to include
                the prediction and linearity losses.
            **kwargs -- additional keyword arguments. Can be used to name the
                Model.
        """
        super().__init__(**kwargs)

        self.n_inputs = n_inputs
        self.n_latent = n_latent
        self.len_time = len_time
        self.num_shifts = num_shifts
        self.num_shifts_middle = num_shifts_middle
        self.outer_encoder = outer_encoder
        self.outer_decoder = outer_decoder

        # Create the inner encoder layer
        self.inner_encoder = keras.layers.Dense(
            n_latent,
            name='inner_encoder',
            activation=None,
            use_bias=False,
            kernel_initializer=identity_init,
            **inner_config)

        # The dynamics matrix, initialized as identity
        self.L = tf.Variable(tf.eye(n_latent), trainable=True)

        # Create the inner decoder layer
        self.inner_decoder = keras.layers.Dense(
            n_inputs,
            name='inner_decoder',
            activation=None,
            use_bias=False,
            kernel_initializer=identity_init,
            **inner_config)

        self.RelMSE = RelMSE(name='RelMSE')
        self.inner_loss_weights = inner_loss_weights
        self.L_diag = L_diag
        self.train_autoencoder_only = train_autoencoder_only

    def call(self, inputs):
        """
        Run given inputs through the newtork.

        Arguments:
            inputs:
            inputs -- A Numpy array or Tensorflow tensor with shape
                (number_of_trajectories, len_time, n_inputs)
            outputs:
            autoencoder_output -- The output of running each time
                step of each trajectory through the autoencoder
            outer_auto_output -- The output of running each time step
                of each trajectory through the outer autoencoder
            predictions -- The predictions for num_steps steps into the
                future for all trajectories and all time steps for which
                num_steps steps into the future are in the data

        """
        # For the shape of the outputs for prediction and linearity loss
        len_pred = self.len_time - self.num_shifts
        len_lin = self.len_time - self.num_shifts_middle

        # Create arrays for prediction and linearity
        pred_inputs = inputs[:, :len_pred, :]
        lin_inputs = inputs[:, :len_lin, :]

        # Inputs for "exact" solutions for prediction and linearity loss
        pred_exact = stack_predictions(inputs, self.num_shifts)
        lin_advanced = stack_predictions(inputs, self.num_shifts_middle)

        #  Reshape inputs as 2D arrays
        auto_inputs, pred_inputs, lin_inputs, lin_advanced = reshape_inputs(
            (inputs, pred_inputs, lin_inputs, lin_advanced))

        # Autoencoder
        partially_encoded = self.outer_encoder(auto_inputs)
        fully_encoded = self.inner_encoder(partially_encoded)
        partially_decoded = self.inner_decoder(fully_encoded)
        autoencoder_output = self.outer_decoder(partially_decoded)

        autoencoder_output = tf.reshape(autoencoder_output,
                                        [-1, self.len_time, self.n_inputs])

        # Outer Autoencoder
        outer_auto_output = self.outer_decoder(partially_encoded)

        outer_auto_output = tf.reshape(outer_auto_output,
                                       [-1, self.len_time, self.n_inputs])

        # Inner Autoencoder Loss
        self.add_loss(self.inner_loss_weights[0]
                      * self.RelMSE(partially_encoded, partially_decoded))

        # If training autoencoder only, output results
        if self.train_autoencoder_only:
            predictions = 0 * pred_exact
            return autoencoder_output, outer_auto_output, predictions

        # Set dynamics matrix L
        if self.L_diag:
            Lmat = tf.linalg.diag(tf.linalg.diag_part(self.L))
        else:
            Lmat = self.L

        # Prediction
        predictions_list = []
        part_encoded_pred = self.outer_encoder(pred_inputs)
        current_encoded = self.inner_encoder(part_encoded_pred)
        for shift in range(self.num_shifts):
            advanced_encoded = tf.matmul(current_encoded, Lmat)
            adv_part_decoded = self.inner_decoder(advanced_encoded)
            advanced_decoded = self.outer_decoder(adv_part_decoded)
            predictions_list.append(tf.reshape(advanced_decoded,
                                               [-1, len_pred, self.n_inputs]))
            current_encoded = tf.identity(advanced_encoded)
        predictions = tf.concat(predictions_list, axis=1)

        # Linearity predictions
        linearity_list = []
        part_encoded_lin = self.outer_encoder(lin_inputs)
        current_encoded = self.inner_encoder(part_encoded_lin)
        for shift in range(self.num_shifts_middle):
            advanced_encoded = tf.matmul(current_encoded, Lmat)
            current_encoded = tf.identity(advanced_encoded)
            linearity_list.append(tf.reshape(current_encoded,
                                             [-1, len_lin, self.n_latent]))
        lin_pred = tf.concat(linearity_list, axis=1)

        lin_part_encoded = self.outer_encoder(lin_advanced)
        lin_exact = self.inner_encoder(lin_part_encoded)
        lin_exact = tf.reshape(
            lin_exact,
            [-1, self.num_shifts_middle * len_lin, self.n_latent])

        # Add Linearity loss
        self.add_loss(self.inner_loss_weights[1]
                      * self.RelMSE(lin_exact, lin_pred))

        return autoencoder_output, outer_auto_output, predictions

def identity_init(shape, dtype=tf.float32):
    """Initialize weight matrices as identity-like matrices."""
    n_rows = shape[0]
    n_cols = shape[1]
    if n_rows >= n_cols:
        A = np.zeros((n_rows, n_cols), dtype=np.float32)
        for col in range(n_cols):
            for row in range(col, col + n_rows - n_cols + 1):
                A[row, col] = 1.0 / (n_rows - n_cols + 1)
    else:
        A = np.zeros((n_rows, n_cols), dtype=np.float32)
        for row in range(n_rows):
            for col in range(row, row + n_cols - n_rows + 1):
                A[row, col] = 1.0 / (n_cols - n_rows + 1)
    return A

def reshape_inputs(inputs):
    """Reshape inputs to be 2D arrays."""
    input_list = []
    for data in inputs:
        input_list.append(tf.reshape(data, [-1, data.shape[-1]]))
    return tuple(input_list)

def stack_predictions(data, num_shifts):
    """Stack inputs for linearity or prediction loss."""
    len_pred = data.shape[1] - num_shifts
    prediction_list = []
    for j in range(num_shifts):
        prediction_list.append(data[:, j + 1:j + 1 + len_pred, :])
    prediction_tensor = tf.concat(prediction_list, axis=1)

    return prediction_tensor


import unittest
import sys
sys.path.append('/python_progs/NewGame 12.11/NewGame 12.11')
from scanner import find_tokens
from scanner import Token
from scanner import print_tokens

# Название файла со сканируемым текстом
F_NAME = "test.txt"
# Токены для сравнения с результатом работы модуля
tokens = []

def new_file(text):
    '''
    Открывает и записывает текст в файл
    Принимает текст
    '''
    file = open(F_NAME, "w")
    file.write(text)
    file.close()

def new_tokens(lis):
    '''
    Добавляет тоекны в список
    Принимает список с параметрами токенов
    '''
    global tokens
    tokens = []
    for line in lis:
        tokens_line = []
        for element in line:
            token = Token(element[0], element[1], element[2], element[3])
            tokens_line += [token]
        tokens += [tokens_line]

def equal_tokens(token1, token2):
    '''
    Функция для сравнения токенов
    Принимает два сравниваемых токена
    Если они одинаковые (их поля) то возвращает True
    иначе возвращает False
    '''
    if token1.type != token2.type:
        return False
    if token1.value != token2.value:
        return False
    if token1.line_num != token2.line_num:
        return False
    if token1.char_num != token2.char_num:
        return False
    return True

def are_equal(list1, list2):
    '''
    Функция для сравнения списка токенов
    Принимает два сравниваемых списка
    Если они одинаковые то возвращает True
    иначе возвращает False
    '''
    if len(list1) != len(list2):
        print("not equal len of lists: len1 = " + str(len(list1)) + " len2 = " + str(len(list2)))
        return False
    for i in range(len(list1)):
        if len(list1[i]) != len(list2[i]):
            print("not equal len[" + i +"] of lists: len1 = " + str(len(list1[i])) + " len2 = " + str(len(list2[i])))
            return False
        for j in range(len(list1[i])):
            if not equal_tokens(list1[i][j], list2[i][j]):
                print("found not equal tokens:\n" + str(list1[i][j]) + "\n" + str(list2[i][j]))
                return False
    return True

class TestScanner(unittest.TestCase):
    '''
    Класс для тестирования сканера
    '''
    def test_comments_symbols(self):
        '''
        Тестирует комментарии и символы
        '''
        print("\ntest_comments_symbols")
        new_file(': spam:: .# . : spam spam spam')
        lis = [[ # line 0
            [1, 1,   "COLON",       None],
            [1, 3,   "IDENTIFIER",  "spam"],
            [1, 7,   "COLON",       None],
            [1, 8,   "COLON",       None],
            [1, 10,   "DOT",         None],
            ]]
        new_tokens(lis)
        found_tokens = find_tokens(F_NAME)
        self.assertTrue(are_equal(found_tokens, tokens))
        print("test_comments_symbols passed")

    def test_string(self):
        '''
        Тестирует строки
        '''
        print("\ntest_string")
        new_file('"str"."str"' + '\n'+
                 '"str"a"str"' + '\n'+
                 '"str"5"str"' + '\n'+
                 '"str""st#r"#none"')
        lis = [[ # line 0
            [1, 1,   "STRING",      "str"],
            [1, 6,   "DOT",         None],
            [1, 7,   "STRING",      "str"],
            ],[ # line 1
            [2, 1,   "STRING",      "str"],
            [2, 6,   "IDENTIFIER",  "a"],
            [2, 7,   "STRING",      "str"],
            ],[ # line 2
            [3, 1,   "STRING",      "str"],
            [3, 6,   "NUMBER",      5],
            [3, 7,   "STRING",      "str"],
            ],[ # line 3
            [4, 1,   "STRING",      "str"],
            [4, 6,   "STRING",      "st#r"],
            ]]
        new_tokens(lis)
        found_tokens = find_tokens(F_NAME)
        self.assertTrue(are_equal(found_tokens, tokens))
        print("test_string passed")

    def test_words_numbers(self):
        '''
        Тестирует слова и цифры
        '''
        print("\ntest_words_numbers")
        new_file('no no  r2d2'       + '\n'+
                 'room   end '       + '\n'+
                 'height width'      + '\n'+
                 'left right up down'+ '\n'+
                 'img  1 22.333')
        lis = [[ # line 0
            [1, 1,   "IDENTIFIER", "no"],
            [1, 4,   "IDENTIFIER", "no"],
            [1, 8,   "IDENTIFIER", "r2d2"],
            ],[ # line 1
            [2, 1,   "ROOM" , None],
            [2, 8,   "END"  , None],
            ],[ # line 2
            [3, 1,   "HEIGHT" , None],
            [3, 8,   "WIDTH"  , None],
            ],[ # line 3
            [4, 1,   "LEFT" , None],
            [4, 6,   "RIGHT", None],
            [4, 12,  "UP"   , None],
            [4, 15,  "DOWN" , None],
            ],[ # line 4
            [5, 1,   "IMG"      , None],
            [5, 6,   "NUMBER"   , 1],
            [5, 8,   "NUMBER"   , 22],
            [5, 10,   "DOT"      , None],
            [5, 11,  "NUMBER"   , 333],
            ]]
        new_tokens(lis)
        found_tokens = find_tokens(F_NAME)
        self.assertTrue(are_equal(found_tokens, tokens))
        print("test_words_numbers passed")

    def test_error_wrong_symbol(self):
        '''
        Тестирует на ошибку неправильного символа
        '''
        print("\ntest_error_wrong_symbol")
        new_file('not4 5year old')
        print("should be wrong symbol error: line 1, char 6")
        found_tokens = find_tokens(F_NAME)
        self.assertEqual(found_tokens, None)
        print("test_error_wrong_symbol passed")

    def test_error_string(self):
        '''
        Тестирует на ошибку строки
        '''
        print("\ntest_error_string")
        new_file('not4"me')
        print("should be string error: line 1, char 7")
        found_tokens = find_tokens(F_NAME)
        self.assertEqual(found_tokens, None)
        print("test_error_string passed")

if __name__ == "__main__":
    print ("run scanner tests\n")
    unittest.main()
    input()


class Account():
    def __init__(self, owner, balance):
        self.owner = owner
        self.balance = balance

    def deposit(self, balance):
        self.balance += balance
        print("Deposit accepted")

    def withdraw(self, amount):
        if amount <= self.balance:
            self.balance -= amount
            print("Withdrawal accepted")
        else:
            print("Funds Unavailable")

    def __str__(self):
        return f"Account owner: {' ':3}{self.owner}\nAccount balance: {' ':1}${self.balance}"

acct1 = Account("Jose", 100)
print(acct1)

print(acct1.owner)
print(acct1.balance)

acct1.deposit(50)
acct1.withdraw(75)
acct1.withdraw(500)


from typing import Tuple

import numpy as np

from core.filters import DataFilter
from core.filters.kalman import ExtendedKalmanFilter
from core.localization import Localization

__author__ = 'Xomak'

class EkfLocalization(Localization):
    def init_filter(self) -> Tuple[DataFilter, np.array]:
        initial_state = self.robot_model.initial_state
        data_filter = ExtendedKalmanFilter(initial_state.shape[0], self.robot_model.observations_dimension)
        data_filter.set_initial(initial_state, self.robot_model.initial_covariance)
        return data_filter, initial_state

    @property
    def filtered_name(self):
        return "EK filter"

    def get_next_state(self, data_filter: DataFilter, step_data):
        state, covariance = data_filter.predict_update(step_data.transition_function,
                                                       step_data.observations_function,
                                                       step_data.control,
                                                       step_data.observations,
                                                       step_data.process_covariance,
                                                       step_data.observations_covariance)
        return state

loc = EkfLocalization()
loc.visualize()


from __future__ import division
import onmt
import onmt.Markdown
import onmt.modules
import argparse
import torch
import time
import datetime
from onmt.train_utils.trainer import XETrainer
from onmt.modules.Loss import NMTLossFunc, NMTAndCTCLossFunc
from onmt.ModelConstructor import build_model
from options import make_parser
from collections import defaultdict

print("cuda is available: ", torch.cuda.is_available())

parser = argparse.ArgumentParser(description='train.py')
onmt.Markdown.add_md_help_argument(parser)

# Please look at the options file to see the options regarding models and data
parser = make_parser(parser)

opt = parser.parse_args()
print(opt)

# An ugly hack to have weight norm on / off
onmt.Constants.weight_norm = opt.weight_norm
onmt.Constants.checkpointing = opt.checkpointing
onmt.Constants.max_position_length = opt.max_position_length

# Use static dropout if checkpointing > 0
if opt.checkpointing > 0:
    onmt.Constants.static = True
if torch.cuda.is_available() and not opt.gpus:
    print("WARNING: You have a CUDA device, should run with -gpus 0")

torch.manual_seed(opt.seed)

def main():

    if opt.data_format == 'raw':
        start = time.time()

        if opt.data.endswith(".train.pt"):
            print("Loading data from '%s'" % opt.data)
            dataset = torch.load(opt.data)
        else:
            print("Loading data from %s" % opt.data + ".train.pt")
            dataset = torch.load(opt.data + ".train.pt")

        elapse = str(datetime.timedelta(seconds=int(time.time() - start)))
        print("Done after %s" % elapse)

        # For backward compatibility
        train_dict = defaultdict(lambda: None, dataset['train'])
        valid_dict = defaultdict(lambda: None, dataset['valid'])

        train_data = onmt.Dataset(train_dict['src'], train_dict['tgt'],
                                  train_dict['src_atbs'], train_dict['tgt_atbs'],
                                  batch_size_words=opt.batch_size_words,
                                  data_type=dataset.get("type", "text"),
                                  batch_size_sents=opt.batch_size_sents,
                                  multiplier=opt.batch_size_multiplier,
                                  augment=opt.augment_speech,
                                  upsampling=opt.upsampling)
        valid_data = onmt.Dataset(valid_dict['src'], valid_dict['tgt'],
                                  valid_dict['src_atbs'], valid_dict['tgt_atbs'],
                                  batch_size_words=opt.batch_size_words,
                                  data_type=dataset.get("type", "text"),
                                  batch_size_sents=opt.batch_size_sents,
                                  upsampling=opt.upsampling)

        dicts = dataset['dicts']

        print(' * number of training sentences. %d' % len(dataset['train']['src']))
        print(' * maximum batch size (words per batch). %d' % opt.batch_size_words)

    elif opt.data_format == 'bin':
        print("Loading memory binned data files ....")
        start = time.time()
        from onmt.data_utils.IndexedDataset import IndexedInMemoryDataset

        dicts = torch.load(opt.data + ".dict.pt")

        train_path = opt.data + '.train'
        train_src = IndexedInMemoryDataset(train_path + '.src')
        train_tgt = IndexedInMemoryDataset(train_path + '.tgt')

        train_data = onmt.Dataset(train_src,
                                  train_tgt,
                                  batch_size_words=opt.batch_size_words,
                                  data_type="text",
                                  batch_size_sents=opt.batch_size_sents,
                                  multiplier=opt.batch_size_multiplier)

        valid_path = opt.data + '.valid'
        valid_src = IndexedInMemoryDataset(valid_path + '.src')
        valid_tgt = IndexedInMemoryDataset(valid_path + '.tgt')

        valid_data = onmt.Dataset(valid_src,
                                  valid_tgt,
                                  batch_size_words=opt.batch_size_words,
                                  data_type="text",
                                  batch_size_sents=opt.batch_size_sents)

        elapse = str(datetime.timedelta(seconds=int(time.time() - start)))
        print("Done after %s" % elapse)
    elif opt.data_format == 'mmem':
        print("Loading memory mapped data files ....")
        start = time.time()
        from onmt.data_utils.MMapIndexedDataset import MMapIndexedDataset

        dicts = torch.load(opt.data + ".dict.pt")

        train_path = opt.data + '.train'
        train_src = MMapIndexedDataset(train_path + '.src')
        train_tgt = MMapIndexedDataset(train_path + '.tgt')

        train_data = onmt.Dataset(train_src,
                                  train_tgt,
                                  batch_size_words=opt.batch_size_words,
                                  data_type="text",
                                  batch_size_sents=opt.batch_size_sents,
                                  multiplier=opt.batch_size_multiplier)

        valid_path = opt.data + '.valid'
        valid_src = MMapIndexedDataset(valid_path + '.src')
        valid_tgt = MMapIndexedDataset(valid_path + '.tgt')

        valid_data = onmt.Dataset(valid_src,
                                  valid_tgt,
                                  batch_size_words=opt.batch_size_words,
                                  data_type="text",
                                  batch_size_sents=opt.batch_size_sents)
        elapse = str(datetime.timedelta(seconds=int(time.time() - start)))
        print("Done after %s" % elapse)

    else:
        raise NotImplementedError

    additional_data = []
    if opt.additional_data != "none":
        add_data = opt.additional_data.split(";")
        add_format = opt.additional_data_format.split(";")
        assert(len(add_data) == len(add_format))
        for i in range(len(add_data)):
            if add_format[i] == 'raw':
                if add_data[i].endswith(".train.pt"):
                    print("Loading data from '%s'" % opt.data)
                    add_dataset = torch.load(add_data[i])
                else:
                    print("Loading data from %s" % opt.data + ".train.pt")
                    add_dataset = torch.load(add_data[i] + ".train.pt")

                additional_data.append(onmt.Dataset(add_dataset['train']['src'],
                                          dataset['train']['tgt'], batch_size_words=opt.batch_size_words,
                                          data_type=dataset.get("type", "text"),
                                          batch_size_sents=opt.batch_size_sents,
                                          multiplier=opt.batch_size_multiplier,
                                          reshape_speech=opt.reshape_speech,
                                          augment=opt.augment_speech))
            elif add_format[i] == 'bin':

                from onmt.data_utils.IndexedDataset import IndexedInMemoryDataset

                train_path = add_data[i] + '.train'
                train_src = IndexedInMemoryDataset(train_path + '.src')
                train_tgt = IndexedInMemoryDataset(train_path + '.tgt')

                additional_data.append(onmt.Dataset(train_src,
                                       train_tgt,
                                       batch_size_words=opt.batch_size_words,
                                       data_type=opt.encoder_type,
                                       batch_size_sents=opt.batch_size_sents,
                                       multiplier=opt.batch_size_multiplier))

    if opt.load_from:
        checkpoint = torch.load(opt.load_from, map_location=lambda storage, loc: storage)
        print("* Loading dictionaries from the checkpoint")
        dicts = checkpoint['dicts']
    else:
        dicts['tgt'].patch(opt.patch_vocab_multiplier)
        checkpoint = None

    if "src" in dicts:
        print(' * vocabulary size. source = %d; target = %d' %
              (dicts['src'].size(), dicts['tgt'].size()))
    else:
        print(' * vocabulary size. target = %d' %
              (dicts['tgt'].size()))

    print('Building model...')

    model = build_model(opt, dicts)

    if opt.get_context_emb:
        assert opt.enc_pretrained_model == "transformer"
        assert opt.dec_pretrained_model == "transformer"
        print("We use pretrained model to get contextualized embeddings and feed them to the original transformer")

        if opt.get_context_emb == "bert":
            from pretrain_module.configuration_bert import BertConfig
            from pretrain_module.modeling_bert import BertModel
            emb_pretrain_config = BertConfig.from_json_file(opt.emb_pretrained_config_dir + "/" + opt.emb_config_name)
            pretrain_emb = BertModel(emb_pretrain_config,
                                     bert_word_dropout=opt.emb_pretrain_word_dropout,
                                     bert_emb_dropout=opt.emb_pretrain_emb_dropout,
                                     bert_atten_dropout=opt.emb_pretrain_attn_dropout,
                                     bert_hidden_dropout=opt.emb_pretrain_hidden_dropout,
                                     bert_hidden_size=opt.emb_pretrain_hidden_size,
                                     is_decoder=False,
                                     )
            emb_state_dict_file = opt.emb_pretrained_config_dir + "/" + opt.emb_pretrained_state_dict
            emb_model_state_dict = torch.load(emb_state_dict_file, map_location="cpu")
            print("After builing pretrained model we load the state from:\n", emb_state_dict_file)
            pretrain_emb.from_pretrained(pretrained_model_name_or_path=opt.enc_pretrained_config_dir,
                                         model=pretrain_emb,
                                         output_loading_info=True,
                                         state_dict=emb_model_state_dict,
                                         model_prefix=opt.get_context_emb
                                         )
            model.add_module("pretrain_emb", pretrain_emb)
        else:
            print("Warning: contextualized embeddings can be only got from bert or roberta")
            exit(-1)
    print(model)

    """ Building the loss function """

    loss_function = NMTLossFunc(opt.model_size, dicts['tgt'].size(),
                                label_smoothing=opt.label_smoothing,
                                mirror=opt.mirror_loss,
                                fast_xentropy=opt.fast_xentropy
                                )

    n_params = sum([p.nelement() for p in model.parameters()])
    print('* number of all parameters: %d' % n_params)

    n_params_grad = sum([p.nelement() for p in model.parameters() if p.requires_grad==True])
    print('* number of all parameters that need gradient: %d' % n_params_grad)

    n_params_nograd = sum([p.nelement() for p in model.parameters() if p.requires_grad==False])
    print('* number of all parameters that do not need gradient: %d' % n_params_nograd)

    assert n_params == (n_params_grad + n_params_nograd)

    if len(opt.gpus) > 1 or opt.virtual_gpu > 1:
        raise NotImplementedError("Warning! Multi-GPU training is not fully tested and potential bugs can happen.")
    else:
        trainer = XETrainer(model, loss_function, train_data, valid_data, dicts, opt, setup_optimizer=True)
        if len(additional_data) > 0:
            trainer.add_additional_data(additional_data, opt.data_ratio)

    trainer.run(checkpoint=checkpoint)

if __name__ == "__main__":
    main()


from pytube import YouTube, Playlist
from threading import Thread

def main():
    format = input('1 - file\n2 - playlist\n')

    if format == 1:
        with open('link.txt', 'r') as file:
            links = file.text
            print(links)
    else:
        link_playlist = input('Enter a playlist link: ')
        playlist = Playlist(link_playlist)
        links = playlist.video_urls

    for index, link in enumerate(links):
        index += 1
        print(f'Total video --- {len(links)}\n')
        Thread(target=download_video, args=(link, index, )).start()

def download_video(link, number):
    video = YouTube(link)
    stream = video.streams.get_highest_resolution()
    stream.download('video/')
    print('\t' + str(number) + ' --- finish')

if __name__ == '__main__':
    main()


import re
import statistics

_fhandler = open('mbox-short.txt')
count = []
for line in _fhandler:
    line = line.rstrip()
    query = re.findall('^New .*: ([0-9]+)', line)
    if len(query) > 0:
        for number in query:
            count.append(float(number))

print(statistics.mean(count))


from django import forms
from .models import Exercise

WEIGHT_TYPE_CHOICES = [
    ('lb', 'lb'),
    ('kg', 'kg'),
]

class ExerciseForm(forms.ModelForm):
    exercise_text = forms.CharField(widget=forms.TextInput(attrs={'placeholder': 'Exercise',"size": "10"}), required=True)
    class Meta:
        model = Exercise
        fields = ('exercise_text',)


# Standard libaries
import dataclasses
import typing
from typing import Optional, List

# External libaries
import numpy as np

# Internal libaries
import pyscf
import pyscf.cc

# Local modules
import vayesta
from vayesta.core.util import deprecated, dot, einsum, energy_string, getattr_recursive, hstack, log_method, log_time
from vayesta.core.qemb import Fragment as BaseFragment
from vayesta.core.fragmentation import IAO_Fragmentation
from vayesta.core.types import RFCI_WaveFunction, RCCSDTQ_WaveFunction, UCCSDTQ_WaveFunction
from vayesta.core.bath import DMET_Bath
from vayesta.mpi import mpi

from vayesta.ewf import ewf

# Get MPI rank of fragment
get_fragment_mpi_rank = lambda *args: args[0].mpi_rank

@dataclasses.dataclass
class Options(BaseFragment.Options):
    # Inherited from Embedding
    # ------------------------
    t_as_lambda: bool = None  # If True, use T-amplitudes inplace of Lambda-amplitudes
    bsse_correction: bool = None
    bsse_rmax: float = None
    sc_mode: int = None
    nelectron_target: float = (
        None  # If set, adjust bath chemical potential until electron number in fragment equals nelectron_target
    )
    nelectron_target_atol: float = 1e-6
    nelectron_target_rtol: float = 1e-6
    # Calculation modes
    calc_e_wf_corr: bool = None
    calc_e_dm_corr: bool = None
    store_wf_type: str = None  # If set, fragment WFs will be converted to the respective type, before storing them
    # Fragment specific
    # -----------------
    wf_factor: Optional[int] = None
    # TODO: move these:
    # CAS methods
    c_cas_occ: np.ndarray = None
    c_cas_vir: np.ndarray = None
    # --- Solver options
    # "TCCSD-solver":
    tcc_fci_opts: dict = dataclasses.field(default_factory=dict)
    # --- Couple embedding problems (currently only CCSD and MPI)
    # coupled_iterations: bool = None # Now accessible through solver=coupledCCSD setting.

class Fragment(BaseFragment):
    Options = Options

    @dataclasses.dataclass
    class Flags(BaseFragment.Flags):
        # Tailoring and external correction of CCSD
        external_corrections: Optional[List[typing.Any]] = dataclasses.field(default_factory=list)
        # Whether to perform additional checks on external corrections
        test_extcorr: bool = False

    @dataclasses.dataclass
    class Results(BaseFragment.Results):
        e_corr_dm2cumulant: float = None
        n_active: int = None
        ip_energy: np.ndarray = None
        ea_energy: np.ndarray = None
        moms: tuple = None

        @property
        def dm1(self):
            """Cluster 1DM"""
            return self.wf.make_rdm1()

        @property
        def dm2(self):
            """Cluster 2DM"""
            return self.wf.make_rdm2()

    def __init__(self, *args, **kwargs):
        """
        Parameters
        ----------
        base : EWF
            Base EWF object.
        fid : int
            Unique ID of fragment.
        name :
            Name of fragment.
        """
        super().__init__(*args, **kwargs)
        # For self-consistent mode
        self.solver_results = None

    def _reset(self, *args, **kwargs):
        super()._reset(*args, **kwargs)
        # Need to unset these so can be regenerated each iteration.
        self.opts.c_cas_occ = self.opts.c_cas_vir = None

    def set_cas(self, iaos=None, c_occ=None, c_vir=None, minao="auto", dmet_threshold=None):
        """Set complete active space for tailored CCSD and active-space CC methods."""
        if dmet_threshold is None:
            dmet_threshold = 2 * self.opts.bath_options["dmet_threshold"]
        if iaos is not None:
            # Create new IAO fragmentation
            frag = IAO_Fragmentation(self.base, minao=minao)
            frag.kernel()
            # Get IAO and environment coefficients from fragmentation
            indices = frag.get_orbital_fragment_indices(iaos)[1]
            c_iao = frag.get_frag_coeff(indices)
            c_env = frag.get_env_coeff(indices)
            bath = DMET_Bath(self, dmet_threshold=dmet_threshold)
            c_dmet = bath.make_dmet_bath(c_env)[0]
            tol = self.opts.bath_options["occupation_tolerance"]
            c_iao_occ, c_iao_vir = self.diagonalize_cluster_dm(c_iao, c_dmet, tol=2 * tol)
        else:
            c_iao_occ = c_iao_vir = None

        c_cas_occ = hstack(c_occ, c_iao_occ)
        c_cas_vir = hstack(c_vir, c_iao_vir)
        self.opts.c_cas_occ = c_cas_occ
        self.opts.c_cas_vir = c_cas_vir
        return c_cas_occ, c_cas_vir

    def add_external_corrections(
        self, fragments, correction_type="tailor", projectors=1, test_extcorr=False, low_level_coul=True
    ):
        """Add tailoring or external correction from other fragment solutions to CCSD solver.

        Parameters
        ----------
        fragments: list
            List of solved or auxiliary fragments, used for the correction.
        correction_type: str, optional
            Type of correction:
                'tailor': replace CCSD T1 and T2 amplitudes with FCI amplitudes.
                'delta-tailor': Add the difference of FCI and CCSD T1 and T2 amplitudes
                'external': externally correct CCSD T1 and T2 amplitudes from FCI T3 and T4 amplitudes.
            Default: 'tailor'.
        projectors: int, optional
            Maximum number of projections applied to the occupied dimensions of the amplitude corrections.
            Default: 1.
        test_extcorr: bool, optional
            Whether to perform additional checks on the external corrections.
        low_level_coul: bool, optional
            This is an option specific to the 'external' correction.
            If True, then the T3V term is contracted with integrals spanning the 'low-level' (i.e. CCSD) solver, i.e. the cluster being constrained.
            If False, then the T3V term is contracted with the integrals in the 'high-level' (i.e. FCI) solver, i.e. the cluster providing the constraints.
            In general, there should be a slight speed increase, and slight loss of accuracy for the low_level_coul=False option, but in practice, we find only
            minor differences.
            Default: True
        """
        if correction_type not in ("tailor", "delta-tailor", "external"):
            raise ValueError
        if self.solver == "CCSD":
            # Automatically update this cluster to use external correction solver.
            self.solver = "extCCSD"
            self.check_solver(self.solver)
        if self.solver != "extCCSD":
            raise RuntimeError
        if (not low_level_coul) and correction_type != "external":
            raise ValueError(
                "low_level_coul optional argument only meaningful with 'external' correction of fragments."
            )
        if np.any([(getattr_recursive(f, "results.wf", None) is None and not f.opts.auxiliary) for f in fragments]):
            raise ValueError(
                "Fragments for external correction need to be already solved or defined as auxiliary fragments."
            )
        self.flags.external_corrections.extend([(f.id, correction_type, projectors, low_level_coul) for f in fragments])
        self.flags.test_extcorr = test_extcorr

    def clear_external_corrections(self):
        """Remove all tailoring or external correction which were added via add_external_corrections."""
        self.flags.external_corrections = []
        self.flags.test_extcorr = False

    def get_init_guess(self, init_guess, solver, cluster):
        # FIXME
        return {}

    def kernel(self, solver=None, init_guess=None):
        solver = solver or self.solver
        self.check_solver(solver)
        if self.cluster is None:
            raise RuntimeError
        cluster = self.cluster

        if solver == "HF":
            return None

        init_guess = self.get_init_guess(init_guess, solver, cluster)

        # Create solver object
        cluster_solver = self.get_solver(solver)
        # Calculate cluster energy at the level of RPA.
        e_corr_rpa = self.get_local_rpa_correction(cluster_solver.hamil)
        # --- Chemical potential
        cpt_frag = self.base.opts.global_frag_chempot
        if self.opts.nelectron_target is not None:
            cluster_solver.optimize_cpt(
                self.opts.nelectron_target,
                c_frag=self.c_proj,
                atol=self.opts.nelectron_target_atol,
                rtol=self.opts.nelectron_target_rtol,
            )
        elif cpt_frag:
            # Add chemical potential to fragment space
            r = self.get_overlap("cluster|frag")
            if self.base.is_rhf:
                p_frag = np.dot(r, r.T)
                cluster_solver.v_ext = cpt_frag * p_frag
            else:
                p_frag = (np.dot(r[0], r[0].T), np.dot(r[1], r[1].T))
                cluster_solver.v_ext = (cpt_frag * p_frag[0], cpt_frag * p_frag[1])

        # --- Coupled fragments.
        # TODO rework this functionality to combine with external corrections/tailoring.
        if solver == "coupledCCSD":
            if not mpi:
                raise RuntimeError("coupled_iterations requires MPI.")
            if len(self.base.fragments) != len(mpi):
                raise RuntimeError("coupled_iterations requires as many MPI processes as there are fragments.")
            cluster_solver.set_coupled_fragments(self.base.fragments)

        # Normal solver
        if not self.base.opts._debug_wf:
            with log_time(self.log.info, ("Time for %s solver:" % solver) + " %s"):
                cluster_solver.kernel()
        # Special debug "solver"
        else:
            if self.base.opts._debug_wf == "random":
                cluster_solver._debug_random_wf()
            else:
                cluster_solver._debug_exact_wf(self.base._debug_wf)

        if solver.lower() == "dump":
            return

        wf = cluster_solver.wf
        # Multiply WF by factor [optional]
        if self.opts.wf_factor is not None:
            wf.multiply(self.opts.wf_factor)
        # Convert WF to different type [optional]
        if self.opts.store_wf_type is not None:
            wf = getattr(wf, "as_%s" % self.opts.store_wf_type.lower())()
        # ---Make T-projected WF
        pwf = wf
        # Projection of FCI wave function is not implemented - convert to CISD
        if isinstance(wf, RFCI_WaveFunction):
            pwf = wf.as_cisd()
        # Projection of CCSDTQ wave function is not implemented - convert to CCSD
        elif isinstance(wf, (RCCSDTQ_WaveFunction, UCCSDTQ_WaveFunction)):
            pwf = wf.as_ccsd()
        proj = self.get_overlap("proj|cluster-occ")
        pwf = pwf.project(proj, inplace=False)

        # Moments

        moms = cluster_solver.hole_moments, cluster_solver.particle_moments

        # --- Add to results data class
        self._results = results = self.Results(
            fid=self.id,
            n_active=cluster.norb_active,
            converged=cluster_solver.converged,
            wf=wf,
            pwf=pwf,
            moms=moms,
            e_corr_rpa=e_corr_rpa,
        )

        self.hamil = cluster_solver.hamil

        # --- Correlation energy contributions
        if self.opts.calc_e_wf_corr:
            ci = wf.as_cisd(c0=1.0)
            ci = ci.project(proj)
            es, ed, results.e_corr = self.get_fragment_energy(ci.c1, ci.c2, hamil=self.hamil)
            self.log.debug(
                "E(S)= %s  E(D)= %s  E(tot)= %s", energy_string(es), energy_string(ed), energy_string(results.e_corr)
            )
        if self.opts.calc_e_dm_corr:
            results.e_corr_dm2cumulant = self.make_fragment_dm2cumulant_energy(hamil=self.hamil)
        return results

    def get_solver_options(self, solver):
        # TODO: fix this mess...
        # Use those values from solver_options, which are not None
        # (conv_tol, max_cycle, solve_lambda,...)
        solver_opts = {key: val for (key, val) in self.opts.solver_options.items() if val is not None}
        pass_through = []
        if "CCSD" in solver.upper():
            pass_through += ["sc_mode", "dm_with_frozen"]
        for attr in pass_through:
            self.log.debugv("Passing fragment option %s to solver.", attr)
            solver_opts[attr] = getattr(self.opts, attr)

        has_actspace = (
            (solver == "TCCSD")
            or ("CCSDt'" in solver)
            or (solver.upper() == "EBCC" and self.opts.solver_options["ansatz"] == "CCSDt'")
        )
        if has_actspace:
            # Set CAS orbitals
            if self.opts.c_cas_occ is None:
                self.log.warning("Occupied CAS orbitals not set. Setting to occupied DMET cluster orbitals.")
                self.opts.c_cas_occ = self._dmet_bath.c_cluster_occ
            if self.opts.c_cas_vir is None:
                self.log.warning("Virtual CAS orbitals not set. Setting to virtual DMET cluster orbitals.")
                self.opts.c_cas_vir = self._dmet_bath.c_cluster_vir
            solver_opts["c_cas_occ"] = self.opts.c_cas_occ
            solver_opts["c_cas_vir"] = self.opts.c_cas_vir
            if solver == "TCCSD":
                solver_opts["tcc_fci_opts"] = self.opts.tcc_fci_opts
        elif solver.upper() == "DUMP":
            solver_opts["filename"] = self.opts.solver_options["dumpfile"]
        solver_opts["external_corrections"] = self.flags.external_corrections
        solver_opts["test_extcorr"] = self.flags.test_extcorr
        return solver_opts

    # --- Expectation values
    # ----------------------

    # --- Energies

    def get_fragment_energy(self, c1, c2, hamil=None, fock=None, c2ba_order="ba", axis1="fragment"):
        """Calculate fragment correlation energy contribution from projected C1, C2.

        Parameters
        ----------
        c1 : (n(occ-CO), n(vir-CO)) array
            Fragment projected C1-amplitudes.
        c2 : (n(occ-CO), n(occ-CO), n(vir-CO), n(vir-CO)) array
            Fragment projected C2-amplitudes.
        hamil : ClusterHamiltonian object.
            Object representing cluster hamiltonian, possibly including cached ERIs.
        fock : (n(AO), n(AO)) array, optional
            Fock matrix in AO representation. If None, self.base.get_fock_for_energy()
            is used. Default: None.

        Returns
        -------
        e_singles : float
            Fragment correlation energy contribution from single excitations.
        e_doubles : float
            Fragment correlation energy contribution from double excitations.
        e_corr : float
            Total fragment correlation energy contribution.
        """
        if axis1 == "fragment":
            px = self.get_overlap("proj|cluster-occ")

        # --- Singles energy (zero for HF-reference)
        if c1 is not None:
            if fock is None:
                fock = self.base.get_fock_for_energy()
            fov = dot(self.cluster.c_active_occ.T, fock, self.cluster.c_active_vir)
            if axis1 == "fragment":
                e_singles = 2 * einsum("ia,xi,xa->", fov, px, c1)
            else:
                e_singles = 2 * np.sum(fov * c1)
        else:
            e_singles = 0
        # --- Doubles energy
        if hamil is None:
            hamil = self.hamil
        # This automatically either slices a stored ERI tensor or calculates it on the fly.
        g_ovvo = hamil.get_eris_bare(block="ovvo")

        if axis1 == "fragment":
            e_doubles = 2 * einsum("xi,xjab,iabj", px, c2, g_ovvo) - einsum("xi,xjab,ibaj", px, c2, g_ovvo)
        else:
            e_doubles = 2 * einsum("ijab,iabj", c2, g_ovvo) - einsum("ijab,ibaj", c2, g_ovvo)

        e_singles = self.sym_factor * e_singles
        e_doubles = self.sym_factor * e_doubles
        e_corr = e_singles + e_doubles
        return e_singles, e_doubles, e_corr

    # --- Density-matrices

    def _ccsd_amplitudes_for_dm(self, t_as_lambda=False, sym_t2=True):
        wf = self.results.wf.as_ccsd()
        t1, t2 = wf.t1, wf.t2
        pwf = self.results.pwf.restore(sym=sym_t2).as_ccsd()
        t1x, t2x = pwf.t1, pwf.t2
        # Lambda amplitudes
        if t_as_lambda:
            l1, l2 = t1, t2
            l1x, l2x = t1x, t2x
        else:
            l1, l2 = wf.l1, wf.l2
            l1x, l2x = pwf.l1, pwf.l2
        return t1, t2, l1, l2, t1x, t2x, l1x, l2x

    def _get_projected_gamma1_intermediates(self, t_as_lambda=False, sym_t2=True):
        """Intermediates for 1-DM, projected in Lambda-amplitudes and linear T-term."""
        t1, t2, l1, l2, t1x, t2x, l1x, l2x = self._ccsd_amplitudes_for_dm(t_as_lambda=t_as_lambda, sym_t2=sym_t2)
        doo, dov, dvo, dvv = pyscf.cc.ccsd_rdm._gamma1_intermediates(None, t1, t2, l1x, l2x)
        # Correction for term without Lambda amplitude:
        dvo += (t1x - t1).T
        d1 = (doo, dov, dvo, dvv)
        return d1

    def _get_projected_gamma2_intermediates(self, t_as_lambda=False, sym_t2=True):
        """Intermediates for 2-DM, projected in Lambda-amplitudes and linear T-term."""
        t1, t2, l1, l2, t1x, t2x, l1x, l2x = self._ccsd_amplitudes_for_dm(t_as_lambda=t_as_lambda, sym_t2=sym_t2)
        cc = self.mf  # Only attributes stdout, verbose, and max_memory are needed, just use mean-field object
        dovov, *d2rest = pyscf.cc.ccsd_rdm._gamma2_intermediates(cc, t1, t2, l1x, l2x)
        # Correct D2[ovov] part (first element of d2 tuple)
        dtau = (t2x - t2) + einsum("ia,jb->ijab", (t1x - t1), t1)
        dovov += dtau.transpose(0, 2, 1, 3)
        dovov -= dtau.transpose(0, 3, 1, 2) / 2
        d2 = (dovov, *d2rest)
        return d2

    def make_fragment_dm1(self, t_as_lambda=False, sym_t2=True):
        """Currently CCSD only.

        Without mean-field contribution!"""
        d1 = self._get_projected_gamma1_intermediates(t_as_lambda=t_as_lambda, sym_t2=sym_t2)
        dm1 = pyscf.cc.ccsd_rdm._make_rdm1(None, d1, with_frozen=False, with_mf=False)
        return dm1

    def make_fragment_dm2cumulant(
        self, t_as_lambda=False, sym_t2=True, sym_dm2=True, full_shape=True, approx_cumulant=True
    ):
        """Currently MP2/CCSD only"""

        if self.solver == "MP2":
            if approx_cumulant not in (1, True):
                raise NotImplementedError
            t2x = self.results.pwf.restore(sym=sym_t2).as_ccsd().t2
            dovov = 2 * (2 * t2x - t2x.transpose(0, 1, 3, 2)).transpose(0, 2, 1, 3)
            if not full_shape:
                return dovov
            nocc, nvir = dovov.shape[:2]
            norb = nocc + nvir
            dm2 = np.zeros(4 * [norb])
            occ, vir = np.s_[:nocc], np.s_[nocc:]
            dm2[occ, vir, occ, vir] = dovov
            dm2[vir, occ, vir, occ] = dovov.transpose(1, 0, 3, 2)
            return dm2

        cc = d1 = None
        d2 = self._get_projected_gamma2_intermediates(t_as_lambda=t_as_lambda, sym_t2=sym_t2)
        dm2 = pyscf.cc.ccsd_rdm._make_rdm2(cc, d1, d2, with_dm1=False, with_frozen=False)
        if approx_cumulant == 2:
            raise NotImplementedError
        elif approx_cumulant in (1, True):
            pass
        elif not approx_cumulant:
            # Remove dm1(cc)^2
            dm1x = self.make_fragment_dm1(t_as_lambda=t_as_lambda, sym_t2=sym_t2)
            dm1 = self.results.wf.make_rdm1(with_mf=False)
            dm2 -= (
                einsum("ij,kl->ijkl", dm1, dm1x) / 2
                + einsum("ij,kl->ijkl", dm1x, dm1) / 2
                - einsum("ij,kl->iklj", dm1, dm1x) / 4
                - einsum("ij,kl->iklj", dm1x, dm1) / 4
            )

        if sym_dm2 and not sym_t2:
            dm2 = (dm2 + dm2.transpose(1, 0, 3, 2) + dm2.transpose(2, 3, 0, 1) + dm2.transpose(3, 2, 1, 0)) / 4
        return dm2

    # def make_partial_dm1_energy(self, t_as_lambda=False):
    #    dm1 = self.make_partial_dm1(t_as_lambda=t_as_lambda)
    #    c_act = self.cluster.c_active
    #    fock = np.linalg.multi_dot((c_act.T, self.base.get_fock(), c_act))
    #    e_dm1 = einsum('ij,ji->', fock, dm1)
    #    return e_dm1

    @log_method()
    def make_fragment_dm2cumulant_energy(self, hamil=None, t_as_lambda=False, sym_t2=True, approx_cumulant=True):
        if hamil is None:
            hamil = self.hamil

        # This is a refactor of original functionality with three forks.
        #   - MP2 solver so dm2 cumulant is just ovov, and we just want to contract this.
        #   - CCSD solver so want to use approximate cumulant and can use optimal contraction of different ERI blocks
        #     making use of permutational symmetries.
        #   - All other solvers where we just use a dense eri contraction and may or may not use the approximate
        #     cumulant.
        # With the new hamiltonian object we can always use the optimal contraction for the approximate cumulant,
        # regardless of solver, and we support `approx_cumulant=False` for CCSD.

        if self.solver == "MP2":
            # This is just ovov shape in this case. TODO neater way to handle this?
            dm2 = self.make_fragment_dm2cumulant(
                t_as_lambda=t_as_lambda, sym_t2=sym_t2, approx_cumulant=approx_cumulant, full_shape=False
            )
            return 2 * einsum("ijkl,ijkl->", hamil.get_eris_bare("ovov"), dm2) / 2
        elif approx_cumulant:
            # Working hypothesis: this branch will effectively always uses `approx_cumulant=True`.
            eris = hamil.get_dummy_eri_object(force_bare=True, with_vext=False)
            d2 = self._get_projected_gamma2_intermediates(t_as_lambda=t_as_lambda, sym_t2=sym_t2)
            return vayesta.core.ao2mo.helper.contract_dm2intermeds_eris_rhf(d2, eris) / 2
        else:
            dm2 = self.make_fragment_dm2cumulant(
                t_as_lambda=t_as_lambda, sym_t2=sym_t2, approx_cumulant=approx_cumulant, full_shape=True
            )
            e_dm2 = einsum("ijkl,ijkl->", hamil.get_eris_bare(), dm2) / 2
        return e_dm2

    # --- Other
    # ---------

    def get_fragment_bsse(self, rmax=None, nimages=5, unit="A"):
        self.log.info("Counterpoise Calculation")
        self.log.info("************************")
        if rmax is None:
            rmax = self.opts.bsse_rmax

        # Atomic calculation with atomic basis functions:
        # mol = self.mol.copy()
        # atom = mol.atom[self.atoms]
        # self.log.debugv("Keeping atoms %r", atom)
        # mol.atom = atom
        # mol.a = None
        # mol.build(False, False)

        natom0, e_mf0, e_cm0, dm = self.counterpoise_calculation(rmax=0.0, nimages=0)
        assert natom0 == len(self.atoms)
        self.log.debugv("Counterpoise: E(atom)= % 16.8f Ha", e_cm0)

        # natom_list = []
        # e_mf_list = []
        # e_cm_list = []
        r_values = np.hstack((np.arange(1.0, int(rmax) + 1, 1.0), rmax))
        # for r in r_values:
        r = rmax
        natom, e_mf, e_cm, dm = self.counterpoise_calculation(rmax=r, dm0=dm)
        self.log.debugv(
            "Counterpoise: n(atom)= %3d  E(mf)= %16.8f Ha  E(%s)= % 16.8f Ha", natom, e_mf, self.solver, e_cm
        )

        e_bsse = self.sym_factor * (e_cm - e_cm0)
        self.log.debugv("Counterpoise: E(BSSE)= % 16.8f Ha", e_bsse)
        return e_bsse

    def counterpoise_calculation(self, rmax, dm0=None, nimages=5, unit="A"):
        mol = self.make_counterpoise_mol(rmax, nimages=nimages, unit=unit, output="pyscf-cp.txt")
        # Mean-field
        # mf = type(self.mf)(mol)
        mf = pyscf.scf.RHF(mol)
        mf.conv_tol = self.mf.conv_tol
        # if self.mf.with_df is not None:
        #    self.log.debugv("Setting GDF")
        #    self.log.debugv("%s", type(self.mf.with_df))
        #    # ONLY GDF SO FAR!
        # TODO: generalize
        if self.base.kdf is not None:
            auxbasis = self.base.kdf.auxbasis
        elif self.mf.with_df is not None:
            auxbasis = self.mf.with_df.auxbasis
        else:
            auxbasis = None
        if auxbasis:
            mf = mf.density_fit(auxbasis=auxbasis)
        # TODO:
        # use dm0 as starting point
        mf.kernel()
        dm0 = mf.make_rdm1()
        # Embedded calculation with same options
        ecc = ewf.EWF(mf, solver=self.solver, bno_threshold=self.bno_threshold, options=self.base.opts)
        ecc.make_atom_cluster(self.atoms, options=self.opts)
        ecc.kernel()

        return mol.natm, mf.e_tot, ecc.e_tot, dm0


from model import DnCNN
from utils import MSE_loss
import utils as utils

from data_loader import DenoisingDataset
import data_loader as dl

import easydict
import json
import wandb
import os
import numpy as np
import torch
import torch.nn.functional as f
from torch.utils.data import DataLoader
import torch.optim as optim
from torchvision.transforms import ToPILImage

from skimage.io import imread
from skimage.metrics import peak_signal_noise_ratio as psnr
from skimage.metrics import structural_similarity as ssim
from sklearn.model_selection import KFold # cross-validation

# 실험 셋팅
expname = '[name]_exp[num]'
filename = expname+'.json'
args = easydict.EasyDict({"model": 'DnCNN', "batch_size": 128, "train_data": 'data/train', 'sigma': 25, "epoch": 1, "lr": 1e-3, "optimizer":'Adam', "scheduler": 'MultiStepLR'})
with open(filename,'w', encoding='utf-8') as f:
    json.dump(args, f, indent='\t')
# 모델, 배치 사이즈, 학습데이터, 노이즈 레벨, 에포크, 학습율, (옵티마이저, 스케쥴러)
# batch_size -> 64, 128 / sigma -> 10, 25, 50 / optimizer -> adam, sgd / num_layer -> 17 20 25

# wandb (wandb.login 필요)
wandb.init(
    # set the wandb project where this run will be logged
    project="dncnn",
    name=expname,

    # track hyperparameters and run metadata
    config={
    "learning_rate": args.lr,
    "architecture": args.model,
    "epochs": args.epoch,
    "batch_size" : args.batch_size,
    "sigma" : args.sigma,
    "num_layer" : args.num_layer,
    "optimizer" : args.optimizer,
    "lr_scheduler" : args.scheduler
    }
)

# cuda
cuda = torch.cuda.is_available()
# setting
k = 3 # k_fold cross validation
batch_size = args.batch_size
n_epoch = args.epoch
sigma = args.sigma
num_layer = args.num_layer
optimizer = args.optimizer
scheduler = args.scheduler
save_dir = os.path.join('models', args.model+'_' + 'sigma' + str(sigma))
if not os.path.exists(save_dir):
    os.mkdir(save_dir)

# from dataloader_utils import train, train_dataloader, test_dataloader
def train_dataloader(train_dataset):
    train_dataloader = DataLoader(
        dataset=train_dataset,
        batch_size=batch_size,
        num_workers=2,
        persistent_workers=True,
        sampler=torch.utils.data.RandomSampler(train_dataset),
    )
    return train_dataloader

if __name__ == '__main__':
    model = DnCNN()
    model.train()
    criterion = MSE_loss()

    xs = dl.datagenerator(data_dir=args.train_data)
    xs = xs.astype('float32')/255.0
    xs = torch.from_numpy(xs.transpose((0, 3, 1, 2)))
    DDataset = DenoisingDataset(xs, sigma)
    DLoader = DataLoader(dataset=DDataset, num_workers=4, drop_last=True, batch_size=batch_size, shuffle=True)
    kfold = KFold(n_splits=k, shuffle=True, random_state=42)
    for fold, (train_indices, val_indices) in enumerate(kfold.split(DDataset)):
        print(f"fold: {fold+1}")

        train_dataset = torch.utils.data.Subset(DDataset, train_indices)
        val_dataset = torch.utils.data.Subset(DDataset, val_indices)

        train_loader = DataLoader(dataset=train_dataset, num_workers=4, drop_last=True, batch_size=batch_size, shuffle=True)
        val_loader = DataLoader(dataset=val_dataset, num_workers=4, drop_last=True, batch_size=batch_size, shuffle=True)

        model = DnCNN(num_layers=num_layer)

        if cuda:
            model = model.cuda()
        if optimizer == 'Adam':
            optimizer = optim.Adam(model.parameters(), lr=args.lr)
        else:
            optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=0.9)
        scheduler = utils.get_lr_scheduler(args.scheduler, optimizer=optimizer)

        # Measure time in PyTorch
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)

        for epoch in range(0, n_epoch):
            epoch_loss = 0

            # start time
            start.record()
            for n_count, xy in enumerate(DLoader):
                low_img, gt_img = xy[0].cuda(), xy[1].cuda() # low_img: noisy image, gt_img:ground truth
                optimizer.zero_grad()
                pred_img = model(low_img)
                loss = criterion(pred_img, gt_img)
                epoch_loss = loss.item()
                loss.backward()
                optimizer.step()
            # end time
            end.record()
            torch.cuda.synchronize()
            scheduler.step()
            # test
            test_gt_img = np.array(imread('/content/drive/MyDrive/DnCNN/data/Set12/01.png'), dtype=np.float32)/255.0
            np.random.seed(seed=0)
            noisy = test_gt_img + np.random.normal(0, sigma/255.0, test_gt_img.shape)  # Add Gaussian noise without clipping
            noisy = noisy.astype(np.float32)
            tensor_noisy = torch.from_numpy(noisy).view(1, -1, noisy.shape[0], noisy.shape[1])
            tensor_noisy = tensor_noisy.cuda()
            pred_img = model(tensor_noisy)  # inference
            pred_img = pred_img.view(noisy.shape[0], noisy.shape[1])
            test_pred_img = pred_img.cpu().detach().numpy().astype(np.float32)

            psnr_score = psnr(test_gt_img, test_pred_img)
            ssim_score = ssim(test_gt_img.squeeze(), test_pred_img.squeeze())
            wandb.log({'loss' : epoch_loss/batch_size,'time' : start.elapsed_time(end), 'image':wandb.Image(test_pred_img),'test_psnr':psnr_score, 'test_ssim':ssim_score})
            print(f'epoch : {epoch+1}, loss : {epoch_loss/batch_size}, time : {start.elapsed_time(end)}')
        torch.save(model, os.path.join(save_dir, f'model_{epoch+1}_{fold+1}fold.pth'))


import pytest

from zmq_tubes.matcher import TopicMatcher

class TestTopicMatcher:

    @pytest.mark.parametrize("sub,topic", [
        ("foo/bar", "foo/bar"),
        ("foo/+", "foo/bar"),
        ("foo/+/baz", "foo/bar/baz"),
        ("foo/+/#", "foo/bar/baz"),
        ("A/B/+/#", "A/B/B/C"),
        ("#", "foo/bar/baz"),
        ("#", "/foo/bar"),
        ("/#", "/foo/bar"),
        ("$SYS/bar", "$SYS/bar"),
    ])
    def test_matching(self, sub, topic):
        t = TopicMatcher()
        t.set_topic(sub, True)
        assert t.match(topic)

    @pytest.mark.parametrize("sub,topic", [
        ("test/6/#", "test/3"),
        ("foo/bar", "foo"),
        ("foo/+", "foo/bar/baz"),
        ("foo/+/baz", "foo/bar/bar"),
        ("foo/+/#", "fo2/bar/baz"),
        ("/#", "foo/bar"),
        ("#", "$SYS/bar"),
        ("$BOB/bar", "$SYS/bar"),
    ])
    def test_not_matching(self, sub, topic):
        t = TopicMatcher()
        t.set_topic(sub, True)
        assert t.match(topic) is None

    @pytest.mark.parametrize("filter,val", [
        ('#', [('g/p/+/i', 'i'), ('g/p/+/o', 'o'), ('g/s', 's')]),
        ('g/#', [('g/p/+/i', 'i'), ('g/p/+/o', 'o'), ('g/s', 's')]),
        ('g/s', [('g/s', 's')]),
        ('g/+/+/i', [('g/p/+/i', 'i')]),
        ('g/+/+/+', [('g/p/+/i', 'i'), ('g/p/+/o', 'o')]),
        ('g/+', [('g/s', 's')]),
    ])
    def test_filtering(self, filter, val):
        t = TopicMatcher()
        t.set_topic('g/p/+/i', 'i')
        t.set_topic('g/p/+/o', 'o')
        t.set_topic('g/s', 's')
        assert set(t.filter(filter)) == set(val)


"""
 Project : Turtle drawing of a rectangle
 Author : 심상혁
 Date of first create : March 6, 2023
 Major features:
   -draw a rectangle with the width and length entered
 Date of last update : March 8, 2023
 Revised contesnt :
   -v1.0 : March 6, 2023
        * first create
   -v1.1 : March 8, 2023
        * revised the comment
"""
import turtle

# 가로 입력
width_str = input("width of rectangle = ")
width = int(width_str)    # 입력받은 가로를 int형으로 변경
# 세로 입력
length_str = input("length of rectangle = ")
length = int(length_str)    # 입력받은 세로를 int형으로 변경
# 시작좌표(사각형의 왼쪽 위)
start_width = 0 - (width / 2)
start_length = 0 + (length / 2)
# 시작좌표를 기준으로 사각형의 꼭짓점 리스트 생성
rectangle = [(start_width, start_length),   # 왼쪽 위
             (start_width + width, start_length),   # 오른쪽 위
             (start_width + width, start_length - length),    # 오른쪽 아래
             (start_width, start_length - length)]    # 왼쪽 아래
t = turtle.Turtle() # 터틀 객체 생성
t.shape("turtle"); t.color("red"); t.width(5); t.speed(1)
t.up(); t.goto(0, 0); t.down()  # 펜을 들고 0,0으로 이동
# 사각형 꼭짓점 리스트를 이용해 사각형을 그린다
for i in range(len(rectangle)):
    x, y = rectangle[i]
    if i == 0:
        t.up(); t.goto(x, y); t.down()
    else:
        t.goto(x, y)
sx, sy = rectangle[0]
t.goto(sx, sy)
# 다 그린 후 펜을 들고 0,0으로 이동
t.up(); t.goto(0, 0); t.down()

컴퓨팅사고와파이썬프로그래밍/Turtle drawing of a rectangle.py


import numpy as np
import math
import matplotlib.pyplot as plt
import pandas as pd
import argparse
import os
import seaborn as sns
import sys

def printing_results(y_test, y_pred, y_sk, house_list):
    from sklearn.metrics import accuracy_score as sk_a_s
    print(f'\nPrediction sklearn = \n{y_sk}\n')
    for i in range(4):
        print(f'Sklearn Accuracy {house_list[i]}:'
              f'{sk_a_s(np.where(y_test != i + 1, 0, y_test), np.where(y_sk != i + 1, 0, y_sk)) * 100:.2f}%')

    print(f'\n\x1b[1;30;43mSklearn Global accuracy: {sk_a_s(y_test, y_sk) * 100:.2f}% \x1b[0m\n')

def compare_with_sk(x_train, x_test, y_train, y_test, y_pred, house_list, num_labels, num_features):
    from sklearn.linear_model import LogisticRegression
    logreg = LogisticRegression(solver='lbfgs', multi_class='multinomial', max_iter=15000)
    logreg.fit(x_train[:, 1:], y_train)
    y_sk = logreg.predict(x_test[:, 1:])
    printing_results(y_test, y_pred, y_sk, house_list)

    # Save metrics
    os.makedirs('thetas', exist_ok=True)
    dict_thetas = dict()
    for i in range(num_labels):
        list_thetas = [logreg.intercept_[i], ]
        for c in range(num_features):
            list_thetas.append(float(logreg.coef_[i][c]))

        dict_thetas[f'label_{i}'] = list_thetas

    df_coefs = pd.DataFrame([dict_thetas[key] for key in dict_thetas.keys()],
                            index=[f'label_{n}' for n in range(num_labels)],
                            columns=[f'theta_{n}' for n in range(num_features + 1)])

    df_coefs.to_excel(os.path.join('thetas', f'sk_coefs.xlsx'))

def acc_score(true, pred):
    c = 0
    length = len(true)
    for i in range(length):
        if true[i] != pred[i]:
            c += 1
    return 1 - (c / length)

def plot_heat_map(y_test, y_pred):
    from sklearn import metrics
    cnf_matrix = metrics.confusion_matrix(y_test, y_pred)
    class_names = [0, 1]  # name  of classes
    fig, ax = plt.subplots(figsize=[7.5, 5])
    tick_marks = np.arange(len(class_names))
    plt.xticks(tick_marks, class_names)
    plt.yticks(tick_marks, class_names)
    sns.heatmap(pd.DataFrame(cnf_matrix), annot=True, cmap="YlGnBu", fmt='g')
    bottom, top = ax.get_ylim()
    ax.set_ylim(bottom + 0.5, top - 0.5)
    ax.xaxis.set_label_position("top")
    plt.tight_layout()
    plt.title('Confusion matrix', y=1.1)
    plt.ylabel('Actual label')
    plt.xlabel('Predicted label')
    plt.tight_layout()

    if args.save:
        plt.savefig(os.path.join(os.getcwd(), 'prediction_matrix_plot.png'))

    plt.show()

def save_metrics(num_labels, classifiers, list_x, num_features):
    os.makedirs('thetas', exist_ok=True)
    dict_thetas = dict()
    for f in range(num_labels):
        list_thetas = []
        for e, c in enumerate(np.nditer(classifiers[f])):
            list_thetas.append(float(c))

        dict_thetas[f'label_{f}'] = list_thetas

    df_coefs = pd.DataFrame([dict_thetas[key] for key in dict_thetas.keys()],
                            index=[f'label_{n}' for n in range(num_labels)],
                            columns=[list_x[n - 1] if n != 0 else 'theta0' for n in range(num_features + 1)])

    df_coefs.to_excel(os.path.join('thetas', f'coefs.xlsx'))

def sigmoid(z):
    return 1 / (1 + math.e ** (-z))

def cost(theta, X, y):
    predictions = sigmoid(X @ theta)
    predictions[predictions == 1] = 0.999999  # log(1)=0 causes error in division
    predictions[predictions == 0] = 0.000001  # log(0)=1 causes error in prediction
    error = -y * np.log(predictions) - (1 - y) * np.log(1 - predictions)
    return sum(error) / len(y)

def cost_gradient(theta, X, y):
    """
    Makes an hypothesis of y with the actual thetas
    :param theta: List of thetas
    :param X: The features used to make a correlation
    :param y: The y used to confirm or infirm the correlation
    :return: A list of actualised thetas
    """

    # Gets the predictied y with the actual thetas
    predictions = sigmoid(X @ theta)
    return X.transpose() @ (predictions - y) / len(y)

def gradient_descent(X, y, lr, num_features):
    """
    Makes a gradient descent for X to y
    :param X: The features used to make a correlation
    :param y: The y used to confirm or infirm the correlation
    :param lr: The learning rate (pre-defined with others tests)
    :param num_features: Number of features
    :return: A list of thetas
    """

    theta = np.zeros(num_features + 1)
    i = 0
    while (cost(theta, X, y) > 0.15 or i < 35000) and i < 50000:
        theta = theta - (lr * cost_gradient(theta, X, y))
        i += 1

    print(f'Number of iteration: {i}')
    return theta

def compute(x_train, y_train, num_labels, num_features, house_list, list_x):
    """
    Computes our data with a gradient descent to minimize the error and make a file filled with the calculated thetas
    :param x_train: Numpy array of features
    :param y_train: Numpy array of the target variable
    :param num_labels: Number of houses
    :param num_features: Number of features
    :param house_list: List of houses
    :param list_x: List of features
    :return: A classifier (an array containing all thetas for each house)
    """

    # Creates our classifiers to be filled with thetas
    classifiers = np.zeros(shape=(num_labels, num_features + 1))

    # One vs. all algorithm, for each house set all values to 0 except the house we want to get the thetas
    for c in range(1, num_labels + 1):
        print(f'Entering gradient descent for: {house_list[c - 1]}')

        # Set all values to 0 except c
        label = (y_train == c).astype(int)

        # Get the thetas using the gradient descent
        classifiers[(c - 1), :] = gradient_descent(x_train, label, 0.000035, num_features)

    # Save metrics into an .xlsx file
    save_metrics(num_labels, classifiers, list_x, num_features)

    return classifiers

def insert_ones(x):
    """
    A small function that is made to turn a DataFrame onto a numpy array. Insert a new col
     to x and set 1 all over the array
    :param x: a DataFrame
    :return: a numpy array
    """

    # This is the thetas array, so it must be sized: length * number of features + 1. (The + 1 stand for theta0)
    # Theta0 column is going to be multiplied by 1 in the cross product so it wont affect our hypothesis
    X = np.ones(shape=(x.shape[0], x.shape[1] + 1))
    X[:, 1:] = x
    return X

def formatting_train_test(x, y):
    """
    Takes two numpy array to split into samples (train and test) and resizes them
    :param x: Features
    :param y: The target variable
    :return: The samples (x_train, x_test, y_train, y_test)
    """

    # Splits samples train / test
    if args.compare or not args.evaluate:
        if args.compare:
            print(f'\nTraining and comparing model on 100% of the datasert\n')
        x_train, x_test = x, x
        y_train, y_test = y, y
    else:
        print(f'\nTraining model on 80% of the datasert')
        print(f'Evaluating model on 20% of the dataset\n')
        from sklearn.model_selection import train_test_split
        x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=0)

    # Just resize arrays to a one dimensional array
    y_train = np.resize(y_train, y_train.shape[0])
    y_test = np.resize(y_test, y_test.shape[0])

    # Takes a array add a column at start and fills it with 1
    x_train = insert_ones(x_train)
    x_test = insert_ones(x_test)

    return x_train, x_test, y_train, y_test

def formatting_data(df, house_list, list_col):
    """
    Takes a DataFrame and two lists of str, to shape data and defines training and testing samples
     x's and y's
    :param df: A DataFrame
    :param house_list: List of str containing elements to be classified
    :param list_col: list of columns, first the y column others = x
    :return: numpy arrays of training and testing part, and sizes of dataset
    """

    df = df[list_col[:]]
    df = df.dropna().reset_index(drop=True)

    x = df[list_col[1:]]
    x = x.to_numpy()
    y = pd.DataFrame(df['Hogwarts House'])
    for i in range(len(house_list)):
        y['Hogwarts House'] = np.where(y['Hogwarts House'] == house_list[i], i + 1, y['Hogwarts House'])
    y = y.astype('int')
    y = y.to_numpy()

    # Makes train and test samples
    x_train, x_test, y_train, y_test = formatting_train_test(x, y)

    num_features = x.shape[1]
    num_labels = 4

    return x_train, x_test, y_train, y_test, num_features, num_labels

def train(df):
    """
    Global function that take a DataFrame and train a logistic regression out of it
    :param df: A DataFrame
    """

    house_list = ['Ravenclaw', 'Slytherin', 'Gryffindor', 'Hufflepuff']
    list_col = ['Hogwarts House', 'Astronomy', 'Herbology', 'Defense Against the Dark Arts', 'Ancient Runes']

    # Defines samples to train and test if args.compare or args.evaluate
    x_train, x_test, y_train, y_test, num_features, num_labels = formatting_data(df, house_list, list_col)

    # Defines a classifier containing our thetas for each element to classify, i.e. each house
    classifiers = compute(x_train, y_train, num_labels, num_features, house_list, list_col[1:])

    # Defines an array of probabilities, and creates for each value an array to get the four probabilities
    probabilities = sigmoid(x_test @ classifiers.transpose())

    # Get the maximum value of each array in probabilities. The + 1 stands to get the right house (ordered
    # from 0 to 3 but but the prediction expect values from 1 to 4)
    y_pred = probabilities.argmax(axis=1) + 1

    if args.evaluate or args.compare:
        print(f'\ny_test =\n{y_test}\n')
        print(f'\ny_pred =\n{y_pred}\n')

        # Plots a heat map to represent the repartition of right/wrong values for each house
        plot_heat_map(y_test, y_pred)

        for i in range(4):
            print(f'DSLR Accuracy of {house_list[i]}:'
                  f'{acc_score(np.where(y_test != i + 1, 0, y_test), np.where(y_pred != i + 1, 0, y_pred)) * 100:.2f}%')

        print(f'\n\x1b[1;30;42mDSLR Global accuracy: {acc_score(y_test, y_pred) * 100:.2f}% \x1b[0m\n')

        # Makes a comparison of our model vs. Scikit-learn
        if args.compare:
            compare_with_sk(x_train, x_test, y_train, y_test, y_pred, house_list, num_labels, num_features)

def parsing():
    """
    Parses and defines parameters
    :return: _args
    """

    parser = argparse.ArgumentParser(prog='py train.py')
    parser.add_argument('csv_file', help='A csv file containing data')
    parser.add_argument('-c', '--compare', action='store_true', help='Comparative mode', default=False)
    parser.add_argument('-e', '--evaluate', action='store_true', help='Evaluate the model', default=False)
    parser.add_argument('-s', '--save', action='store_true', help='Saving plot', default=False)
    _args = parser.parse_args()

    return _args

if __name__ == '__main__':
    args = parsing()

    file = os.path.join(os.getcwd(), args.csv_file)
    if os.path.exists(file)and os.path.isfile(file) and file.endswith('.csv'):
        df = pd.read_csv(file)
        df = df.dropna()
        if not df.empty:
            train(df)
        else:
            sys.exit(print(f'\x1b[1;37;41mEmpty DataFrame \x1b[0m\n'))
    else:
        sys.exit(print(f'\x1b[1;37;41mThe selected file must be a csv file \x1b[0m\n'))


# -*- coding: utf-8 -*-
''' Home made test.'''

from __future__ import print_function

__author__= "Luis C. Pérez Tato (LCPT) and Ana Ortega (AOO)"
__copyright__= "Copyright 2015, LCPT and AOO"
__license__= "GPL"
__version__= "3.0"
__email__= "l.pereztato@gmail.com"

K= 1000 # Spring constant
l= 100 # Distance between nodes
F= 1 # Force magnitude

import xc
from model import predefined_spaces
from materials import typical_materials

# Model definition
feProblem= xc.FEProblem()
preprocessor=  feProblem.getPreprocessor
nodes= preprocessor.getNodeHandler
# Problem type
modelSpace= predefined_spaces.SolidMechanics2D(nodes)

n1= nodes.newNodeXY(0,0)
n2= nodes.newNodeXY(l,0.0)

# Materials definition
elast= typical_materials.defElasticMaterial(preprocessor, "elast",K)

''' We define nodes at the points where loads will be applied.
    We will not compute stresses so we can use an arbitrary
    cross section of unit area.'''

# Elements definition
elements= preprocessor.getElementHandler
elements.defaultMaterial= elast.name
elements.dimElem= 2 # Dimension of element space
#  sintaxis: Spring[<tag>]
spring= elements.newElement("Spring",xc.ID([n1.tag,n2.tag]))

# Constraints
constraints= preprocessor.getBoundaryCondHandler
#
spc1= constraints.newSPConstraint(n1.tag,0,0.0) # Node 1
spc2= constraints.newSPConstraint(n1.tag,1,0.0)
spc3= constraints.newSPConstraint(n2.tag,1,0.0) # Node 2

# Load case definition.
lp0= modelSpace.newLoadPattern(name= '0')

# we check that loads are cummulated by the way.
lp0.newNodalLoad(n2.tag,xc.Vector([F/2.0,0]))
lp0.newNodalLoad(n2.tag,xc.Vector([F/2.0,0]))
modelSpace.addLoadCaseToDomain(lp0.name) # Append load pattern to domain.
# Solution
result= modelSpace.analyze(calculateNodalReactions= True)

deltax= n2.getDisp[0]
deltay= n2.getDisp[1]
R= n1.getReaction[0]

elements= preprocessor.getElementHandler
spring.getResistingForce()
Ax= spring.getMaterial().getStrain() # Spring elongation

ratio1= (F+R/F)
ratio2= ((K*deltax-F)/F)
ratio3= ((deltax-Ax)/Ax)
'''
print("R= ",R)
print("dx= ",deltax)
print("dy= ",deltay)
print("Ax= ",Ax)
print("ratio1= ",ratio1)
print("ratio2= ",ratio2)
print("ratio3= ",ratio3)
 '''
import os
from misc_utils import log_messages as lmsg
fname= os.path.basename(__file__)
if (abs(ratio1)<1e-5) & (abs(ratio2)<1e-5) & (abs(ratio3)<1e-5):
    print('test '+fname+': ok.')
else:
    lmsg.error(fname+' ERROR.')


import os

from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader

TRAIN_FOLDER = "train"
VAL_FOLDER = "valid"
TEST_FOLDER = "test"

def _assert_dir_exists(filepath):
    if not os.path.exists(filepath):
        raise Exception("Directory '{}' does not exists!".format(filepath))

def _assert_transforms_not_empty(transforms, train_key, val_key, test_key):
    if transforms is None:
        raise Exception("Dictionary 'transforms' cannot be None!")

    if len(transforms.keys()) == 0:
        raise Exception("Dictionary 'transforms' cannot be empty!")

    for transform_key in [train_key, val_key, test_key]:
        if transform_key is None:
            raise Exception("Transform key '{}' cannot be None!".format(transform_key))

        if transform_key not in transforms:
            raise Exception("Transform for key '{}' not found!".format(transform_key))

def load_imagefolders(filepath, transforms,
                      train_key=None, val_key=None, test_key=None):

    train_filepath = os.path.join(filepath, TRAIN_FOLDER)
    val_filepath = os.path.join(filepath, VAL_FOLDER)
    test_filepath = os.path.join(filepath, TEST_FOLDER)

    _assert_dir_exists(train_filepath)
    _assert_dir_exists(val_filepath)
    _assert_dir_exists(test_filepath)
    _assert_transforms_not_empty(transforms, train_key, val_key, test_key)

    train_imagefolder = ImageFolder(train_filepath, transform=transforms[train_key])
    val_imagefolder = ImageFolder(val_filepath, transform=transforms[val_key])
    test_imagefolder = ImageFolder(test_filepath, transform=transforms[test_key])

    return train_imagefolder, val_imagefolder, test_imagefolder

def create_dataloaders(train_imagefolder, val_imagefolder, test_imagefolder,
                       train_batch=32, val_batch=32, test_batch=32,
                       train_shuffle=True, val_shuffle=True, test_shuffle=False):

    train_loader = DataLoader(train_imagefolder, batch_size=train_batch, shuffle=train_shuffle)
    val_loader = DataLoader(val_imagefolder, batch_size=val_batch, shuffle=val_shuffle)
    test_loader = DataLoader(test_imagefolder, batch_size=test_batch, shuffle=test_shuffle)

    return train_loader, val_loader, test_loader


# Calculate the average tips to revenue ratio of the drivers for different pickup locations in sorted format.
from mrjob.job import MRJob
from mrjob.step import MRStep

class LocationwiseTipsToRevenueAnalyzer(MRJob):

    def steps(self):
        return [
            MRStep(mapper=self.mapper, combiner=self.combiner, reducer=self.reducer),
            MRStep(reducer=self.final_sort)
        ]

    def mapper(self, _, line):
        if not line.startswith('VendorID'):
            cols = line.split(',')
            amount = float(cols[16])
            tips = float(cols[13])

            yield int(cols[7]), (amount, tips)

    def combiner(self, pickupLocationId, tripDetails):
        totalTipAmount, totalRevenue = 0, 0
        for t in tripDetails:
            totalRevenue += t[0]
            totalTipAmount += t[1]

        yield pickupLocationId, (totalRevenue, totalTipAmount)

    def reducer(self, pickupLocationId, tripDetails):
        totalTipAmount, totalRevenue = 0, 0
        for revenue, tip in tripDetails:
            totalTipAmount += tip
            totalRevenue += revenue
        ratio = totalTipAmount/totalRevenue

        yield None, (ratio, pickupLocationId)

    def final_sort(self, _, tripDetails):
        sorted_trip_details = sorted(tripDetails, reverse=True)

        for val in sorted_trip_details:
            yield val[1], val[0]

if __name__ == '__main__':
    LocationwiseTipsToRevenueAnalyzer.run()

# python mrtask_e.py input  > mrtask_e_result.txt


import sys
sys.path.append("./")

import torch

class EmpathicStoriesDataset(torch.utils.data.Dataset):
    def __init__(self, task, data_pairs, data_stories, labels):
        self.data_pairs = data_pairs
        self.data_stories = data_stories
        self.task = task
        self.labels = labels

    def __getitem__(self, idx):
        if self.task == "similarity":
            i = self.data_pairs.iloc[idx]
            s1 = i["story_A"].replace("\n", "")
            s2 = i["story_B"].replace("\n", "")
            score = i["similarity_empathy_human_AGG"]
            return [s1, s2, score]

        elif self.task == "summary":
            story = self.data_stories.iloc[idx]["story"].replace("\n", "")
            event = self.data_stories.iloc[idx]["Main Event"]
            emotion = self.data_stories.iloc[idx]["Emotion Description"]
            moral = self.data_stories.iloc[idx]["Moral"]
            combined = ""
            if "EVE" in self.labels:
                combined += "[EVE]" + event
            if "EMO" in self.labels:
                combined += "[EMO]" + emotion
            if "MOR" in self.labels:
                combined += "[MOR]" + moral
            return story, combined

    def __len__(self):
        if self.task == "summary":
            return len(self.stories)
        else:
            return len(self.story_pairs)


#!D:\DEV\Python\Python38-32

import argparse
import logging
import time
from datetime import datetime
from functools import partial
from multiprocessing import Pool, Manager

import xlrd
import xlwt
from xlutils.copy import copy

import scraping

timeout = 7 * 24 * 60 * 60

def main(excel_path):
    try:
        start_time = datetime.now()

        print("存放数据的excel路径: {}".format(excel_path))
        order_length = scraping_and_save_to_excel(excel_path)

        end_time = datetime.now()
        user_time = (end_time - start_time).seconds
        print("共查询数据 {} 条，用时 {}，平均每条用时 {} 秒".format(
            order_length,
            format_datetime(user_time),
            round(user_time / order_length, 2)
        ))
        time.sleep(3)

    except PermissionError as e:
        print_error(e, "请在关闭【order.xls】文件后，重新执行脚本。")
    except BaseException as e:
        print_error(e, "数据处理出错，点击【回车键】关闭命令行")
    finally:
        pass

def print_error(e, text):
    logging.exception(e)
    time.sleep(1)
    input("=======================================\n"
          "==== " + text + " ===\n"
                           "=======================================")

def scraping_and_save_to_excel(excel_path):
    workbook = xlrd.open_workbook(excel_path, formatting_info=True)

    order_num_list = read_order_num_from_excel(workbook)

    # 将每20条数据分组
    group_list = [order_num_list[i:i + 20] for i in range(0, len(order_num_list), 20)]

    num_list_str_list = []
    for i, order_list in enumerate(group_list):
        num_list_iter = map(lambda x: x[1], order_list)
        num_list_str_list.append(",".join(str(v) for v in num_list_iter))

    # 多进程抓取数据
    p = Pool(4)
    wrap_func = partial(scraping.scraping_oder_time, lock=Manager().Lock())
    result_list = p.map(wrap_func, num_list_str_list)
    p.close()
    p.join()

    result_dict = {}
    for x in result_list:
        result_dict.update(x)

    print("查询结果{}条：{}".format(len(result_dict), result_dict))

    # 未获取到数据
    if len(result_dict) != 0:
        write_to_excel(workbook, order_num_list, result_dict, excel_path)

    return len(order_num_list)

# 读取excel单号数据
def read_order_num_from_excel(workbook):
    sheet = workbook.sheet_by_index(0)
    num_list = sheet.col_values(0)
    date_list = sheet.col_values(1)
    filter_order = []
    for num in enumerate(num_list):
        if date_list[num[0]] == '' and num[1] != '':
            filter_order.append(num)
    return filter_order

def write_to_excel(workbook, order_list, result_dict, excel_path):
    copy_workbook = copy(workbook)
    copy_sheet = copy_workbook.get_sheet(0)

    style = get_red_style_cell()

    for j, num in enumerate(order_list):
        row_num = num[0]
        order_num = num[1]
        order_date = result_dict.get(order_num.strip(), -1)
        if order_date == -1:
            continue

        copy_sheet.write(row_num, 1, order_num)
        copy_sheet.write(row_num, 2, order_date['start_time'])
        copy_sheet.write(row_num, 3, order_date['end_time'])

        between_seconds = (
                datetime.strptime(order_date['start_time'], '%Y-%m-%d %H:%M') -
                datetime.strptime(order_date['end_time'], '%Y-%m-%d %H:%M')
        ).seconds

        if between_seconds >= timeout:
            copy_sheet.write(row_num, 4, format_datetime(between_seconds, accurate_unit=1), style)
        else:
            copy_sheet.write(row_num, 4, format_datetime(between_seconds, accurate_unit=1))

        # 设置列宽
        copy_sheet.col(0).width = 5000
        copy_sheet.col(1).width = 5000
        copy_sheet.col(2).width = 5000
        copy_sheet.col(3).width = 5000

        print("第{}行写入数据".format(row_num + 1))
    copy_workbook.save(excel_path)

def get_red_style_cell():
    pattern = xlwt.Pattern()  # Create the Pattern
    pattern.pattern = xlwt.Pattern.SOLID_PATTERN  # May be: NO_PATTERN, SOLID_PATTERN, or 0x00 through 0x12
    pattern.pattern_fore_colour = 2  # May be: 8 through 63. 0 = Black, 1 = White, 2 = Red, 3 = Green, 4 = Blue, 5 = Yellow, 6 = Magenta, 7 = Cyan, 16 = Maroon, 17 = Dark Green, 18 = Dark Blue, 19 = Dark Yellow , almost brown), 20 = Dark Magenta, 21 = Teal, 22 = Light Gray, 23 = Dark Gray, the list goes on...
    style = xlwt.XFStyle()  # Create the Pattern
    style.pattern = pattern  # Add Pattern to Style
    return style

# accurate_unit 0 秒, 1 分, 2 小时, 3 天
def format_datetime(between_seconds, accurate_unit=0):
    if between_seconds == 0:
        return ""

    elif between_seconds < 60:
        return "{} 秒 ".format(between_seconds)

    elif 60 <= between_seconds < 3600:
        return "{} 分 ".format(between_seconds // 60) + (
            format_datetime(between_seconds % 60, accurate_unit) if accurate_unit < 1 else "")

    elif 3600 <= between_seconds < 86400:
        return "{} 小时 ".format(between_seconds // 3600) + (
            format_datetime(between_seconds % 3600, accurate_unit) if accurate_unit < 2 else "")

    elif between_seconds >= 86400:
        return "{} 天 ".format(between_seconds // 86400) + (
            format_datetime(between_seconds % 86400, accurate_unit) if accurate_unit < 3 else "")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="scraping shunfeng order data")
    parser.add_argument("--path", "-p", help="存放数据的excel路径，非必要参数", default="order.xls")
    parser.add_argument("--timeout_day", "-t", help="订单时间差的超时时长，非必要参数", default=7)
    args = parser.parse_args()

    day = args.timeout_day
    timeout = day * 24 * 60 * 60
    main(args.path)