from __future__ import print_function
from builtins import range

import numpy as np
import pandas as pd
import sys
import argparse
import os

from sklearn.datasets import load_boston

from sklearn.linear_model import LinearRegression, Ridge, Lasso, LogisticRegression
from sklearn.svm import SVR
from sklearn.ensemble import RandomForestRegressor

from clr_regressors import KPlaneRegressor, CLRpRegressor, CLRcRegressor, RegressorEnsemble
from evaluate import evaluate_all

def preprocess_data(X, y):
  X -= np.min(X, axis=0, keepdims=True)
  X /= np.max(X, axis=0, keepdims=True) / 2.0
  X -= 1.0
  shuffle_idx = np.random.choice(X.shape[0], X.shape[0], replace=False)
  X = X[shuffle_idx]
  y = y[shuffle_idx]
  return X, y

if __name__ == '__main__':
  parser = argparse.ArgumentParser(description='Search parameters')
  parser.add_argument('--dataset', required=True)
  parser.add_argument('--seed', default=0, type=int)
  parser.add_argument('--n_jobs', default=1, type=int)
  parser.add_argument('--global_parallel', dest='global_parallel', action='store_true')
  args = parser.parse_args()
  np.random.seed(args.seed)

  if args.dataset == 'boston':
    boston = load_boston()
    X = boston.data
    y = boston.target
    constr_id = 8
  elif args.dataset == 'abalone':
    abalone_data = pd.read_csv('data/abalone.data', header=None)
    with open('data/abalone.names', 'r') as fin:
      abalone_descr = fin.read()
    X = pd.get_dummies(abalone_data.iloc[:,:-1], columns=[0]).as_matrix().astype(np.float)
    X = np.hstack((X, (np.digitize(X[:, 2], np.linspace(0.1, 0.2, 10)))[:,np.newaxis]))
    y = abalone_data.iloc[:, 8].as_matrix().astype(np.float)
    constr_id = 10
  elif args.dataset == 'auto-mpg':
    data = pd.read_csv('data/auto-mpg.data', header=None, sep='\s+', na_values='?')
    data = data.dropna()
    X = pd.get_dummies(data.iloc[:,1:-1], columns=[7]).as_matrix().astype(np.float)
    y = data[0].as_matrix().astype(np.float)
    constr_id = 5
  else:
    print("Dataset is not supported")
    sys.exit(0)
  X, y = preprocess_data(X, y)

  params = {
    'lr': [Ridge(alpha=1e-5), X, y],
    'ridge 10.0': [Ridge(alpha=10.0), X, y],
  }

  for C in [0.1, 1.0, 16.0, 32.0, 100.0, 128.0]:
    for g in ['auto', 0.25, 0.5, 1.0]:
      for eps in [2 ** (-8), 0.01, 0.25, 0.5]:
        params['svr C={}, g={}, eps={}'.format(C, g, eps)] = [
          SVR(C=C, gamma=g, epsilon=eps), X, y
        ]
  n_jobs = 1 if args.global_parallel else args.n_jobs
  for max_depth in [None, 10, 50]:
    for max_features in ['auto', 5]:
      for min_samples_split in [2, 10, 30]:
        for min_samples_leaf in [1, 10, 30]:
          params[
            'rf md={}, mf={}, mss={}, msl={}'.format(
              max_depth,max_features, min_samples_split, min_samples_leaf)
          ] = [RandomForestRegressor(
            n_estimators=100, max_depth=max_depth,
            max_features=max_features,
            min_samples_split=min_samples_split,
            min_samples_leaf=min_samples_leaf,n_jobs=n_jobs), X, y]

  # TODO: don't refit the model for weighted/not weighted evaluation
  for f in [True, False]:
    for w in [True, False]:
      for k in [2, 4, 6, 8]:
        for l in [0, 1, 10, 100]:
          params['kplane k={} l={} w={} f={}'.format(k, l, w, f)] = [
            KPlaneRegressor(k, l, weighted=w, fuzzy=f), X, y]
          params['CLR_p k={} l={} w={} f={}'.format(k, l, w, f)] = [
            CLRpRegressor(k, l, weighted=w, fuzzy=f), X, y]
          params['CLR_p LR=1.0 k={} l={} w={} f={}'.format(k, l, w, f)] = [
            CLRpRegressor(k, l, weighted=w, fuzzy=f, clf=LogisticRegression()), X, y]
          params['kplane k={} l={} w={} f={} ens=10'.format(k, l, w, f)] = [
            RegressorEnsemble(KPlaneRegressor(k, l, weighted=w, fuzzy=f)), X, y]
          params['CLR_p k={} l={} w={} f={} ens=10'.format(k, l, w, f)] = [
            RegressorEnsemble(CLRpRegressor(k, l, weighted=w, fuzzy=f)), X, y]
          params['CLR_p LR=1.0 k={} l={} w={} f={} ens=10'.format(k, l, w, f)] = [
            RegressorEnsemble(CLRpRegressor(k, l, weighted=w, fuzzy=f, clf=LogisticRegression())), X, y]

  for k in [2, 4, 6, 8]:
    for l in [0, 1, 10, 100]:
      params['CLR_c k={} l={}'.format(k, l)] = [
        CLRcRegressor(k, l, constr_id=constr_id), X, y]
      params['CLR_c k={} l={} ens=10'.format(k, l)] = [
        RegressorEnsemble(CLRcRegressor(k, l, constr_id=constr_id)), X, y]
      params['kplane k={} l={} w=size'.format(k, l)] = [
        KPlaneRegressor(k, l, weighted='size'), X, y]
      params['kplane k={} l={} w=size ens=10'.format(k, l)] = [
        RegressorEnsemble(KPlaneRegressor(k, l, weighted='size')), X, y]

  results = evaluate_all(
    params,
    file_name="results/{}-tmp1.csv".format(args.dataset),
    n_jobs=args.n_jobs,
    gl_parallel=args.global_parallel,
  )
  results = results.sort_values('test_mse_mean')

  add_params = {}
  # assuming ensembles are always best
  algos = [CLRpRegressor, KPlaneRegressor]
  algo_names = ['CLR_p', 'kplane']
  for algo, algo_name in zip(algos, algo_names):
    for idx in results.index:
      if algo_name in idx:
        k = int(idx.split()[1].split('=')[1])
        l = int(idx.split()[2].split('=')[1])
        w = idx.split()[3].split('=')[1] == 'True'
        f = idx.split()[4].split('=')[1] == 'True'
        for alpha in [0.01, 0.1, 1.0, 10.0, 100.0]:
          add_params[idx + ' Lasso {}'.format(alpha)] = [
            RegressorEnsemble(algo(k, l, weighted=w, fuzzy=f, clr_lr=Lasso(alpha))), X, y
          ]
          add_params[idx + ' Ridge {}'.format(alpha)] = [
            RegressorEnsemble(algo(k, l, weighted=w, fuzzy=f, clr_lr=Ridge(alpha))), X, y
          ]
        break
  for idx in results.index:
    if 'CLR_c' in idx:
      k = int(idx.split()[1].split('=')[1])
      l = int(idx.split()[2].split('=')[1])
      for alpha in [0.01, 0.1, 1.0, 10.0, 100.0]:
        add_params[idx + ' Lasso {}'.format(alpha)] = [
          RegressorEnsemble(CLRcRegressor(k, l, constr_id=constr_id, clr_lr=Lasso(alpha))), X, y
        ]
        add_params[idx + ' Ridge {}'.format(alpha)] = [
          RegressorEnsemble(CLRcRegressor(k, l, constr_id=constr_id, clr_lr=Ridge(alpha))), X, y
        ]
      break
  add_results = evaluate_all(
    add_params,
    file_name="results/{}-tmp2.csv".format(args.dataset),
    n_jobs=args.n_jobs,
    gl_parallel=args.global_parallel,
  )

  res_complete = results.append(add_results)
  res_complete.to_csv("results/{}.csv".format(args.dataset))
  os.remove("results/{}-tmp1.csv".format(args.dataset))
  os.remove("results/{}-tmp2.csv".format(args.dataset))


import numpy as np
import matplotlib.pyplot as plt
from scipy.io import wavfile

def plot_fourier_spectrum(freq, spectrum, title):
    plt.plot(freq, np.abs(spectrum))
    plt.xlabel("Frequency (Hz)")
    plt.ylabel("Magnitude")
    plt.title(title)

# Load the WAV file
sample_rate, audio_data = wavfile.read("C:/Users/barwa/Desktop/Math Project/InputData.wav.wav")

# Extract the mono audio channel (if stereo)
if len(audio_data.shape) > 1:
    audio_data = audio_data[:, 0]

# Perform the DFT
spectrum = np.fft.fft(audio_data)
freq = np.fft.fftfreq(len(audio_data), d=1/sample_rate)

# Reconstruct the signal using only the first five Fourier coefficients
reconstructed_spectrum = np.zeros_like(spectrum)
reconstructed_spectrum[:5] = spectrum[:5]
reconstructed_signals = []
for i in range(5):
    reconstructed_spectrum_i = np.zeros_like(spectrum)
    reconstructed_spectrum_i[i] = spectrum[i]
    reconstructed_signal_i = np.fft.ifft(reconstructed_spectrum_i)
    reconstructed_signals.append(reconstructed_signal_i)

# Plot the original waveform
plt.figure(figsize=(8, 6))
plt.subplot(6, 1, 1)
plt.plot(audio_data)
plt.xlabel("Time (samples)")
plt.ylabel("Amplitude")
plt.title("Original Waveform")

# Plot the reconstructed waveforms
for i in range(5):
    plt.subplot(6, 1, i+2)
    plt.plot(reconstructed_signals[i].real)
    plt.xlabel("Time (samples)")
    plt.ylabel("Amplitude")
    plt.title("Reconstructed Waveform ({} Fourier coefficient)".format(i+1))

plt.tight_layout()
plt.show()

# Plot the Fourier spectra of the reconstructed waveforms
plt.figure(figsize=(8, 6))
for i in range(5):
    plt.subplot(5, 1, i+1)
    reconstructed_spectrum_i = np.zeros_like(spectrum)
    reconstructed_spectrum_i[i] = spectrum[i]
    reconstructed_signal_i = np.fft.ifft(reconstructed_spectrum_i)
    reconstructed_spectrum_i = np.fft.fft(reconstructed_signal_i)
    plot_fourier_spectrum(freq, reconstructed_spectrum_i, "Frequency Spectrum ({} Fourier coefficient)".format(i+1))

plt.tight_layout()
plt.show()


import requests
from html import unescape
from bs4 import BeautifulSoup
from tkinter import *
from tkinter import ttk
import os

def review_button():
    tp = Toplevel()
    tp.title("Flashcard_Studio/Sets")
    tp.geometry("800x350")
    global setnames
    global setnamesStrVar
    setnames = []
    label2 = Label(tp, text = "Sets:")
    label2.pack()
    # listdir collects set names to be displayed in OptionMenu() and which when one is selected and the
    # setselbut Button is clicked will run selectsetoptions()
    for path in os.listdir("flashdata/"):
        setnames.append(path)
    setnamesStrVar = StringVar()
    setnamesStrVar.set(setnames[0])
    dropdownsetsel = ttk.OptionMenu(tp, setnamesStrVar, *setnames)
    setselbut = Button(tp, text = "Set Options:", command = selectsetoptions)
    dropdownsetsel.pack()
    setselbut.pack()

def selectsetoptions():
    global selectsetdata
    filenames = setnamesStrVar.get()
    # Gets the name of the set picked to be edited
    with open(f"flashdata/{filenames}", "r") as file:
        selectsetdata = []
        file.seek(0)
        selectsetdata = file.read().splitlines()
        file.close()
    # Opens the slected set file and puts all the data into a 1D list formatted [word, defintion]
    tpp = Toplevel()
    tpp.title("Flashcard_Studio/Sets/Mode")
    flashcardmodebutton = Button(tpp, text = "Flashcards", command = flashcardmode)
    flashcardmodebutton.pack()
    # Currently there is only one option for flashcards available, and that is flashcardmode()

def flashcardmode():
    global dataindex
    global flipped
    global flashcardflip
    flipped = False
    dataindex = 0
    fm = Toplevel()
    fm.geometry("1200x625")
    fm.title("Flashcard_Studio/Sets/Mode/Flashcards")
    flashcardflip = Button(fm, text = selectsetdata[dataindex], height = 10, width = 120, command = flip)
    # The flashcards are just a massive button that uses flip() to change the text on the button.
    prevflashbutton = Button(fm, text = "<", command = prevflash)
    nextflashbutton = Button(fm, text = ">", command = nextflash)
    flashcardflip.pack()
    prevflashbutton.pack()
    nextflashbutton.pack()

def flip():
    global flipped
    if flipped == False:
        flashcardflip.config(text = selectsetdata[dataindex+1])
        # Because of the way the list of flashcards is formatted ([word, definition]), the
        # flashcards are in groups of 2.
        flipped = True
    else:
        flashcardflip.config(text = selectsetdata[dataindex])
        flipped = False

def prevflash():
    global dataindex
    dataindex = dataindex - 2
    print(dataindex)

def nextflash():
    global dataindex
    dataindex = dataindex + 2
    print(dataindex)

def scrape():
    HEADERS = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/87.0.4280.88 Safari/537.36"
    }
    soup = BeautifulSoup(requests.get(entry_text.get(), headers=HEADERS).text, "html.parser")
    # Gets the actual HTML file of the Quizlet that you would like to scrape.
    flashcard_title = soup.find("h1").text

    file = open(f'flashdata/{flashcard_title}.txt', 'w')
    # Creates and opens a new file or opens an old one ready for re/writing

    setdata = soup.find_all("div", class_ = "SetPageTerms-term")
    setdata += soup.find_all("div", class_ = "SetPageTerms-term SetPageTerms-term--beyondSignupThreshold")
    # Collecting the setdata divs, not the raw data.

    for index, flashcards in enumerate(setdata):
        flashlist = setdata[index]
        flashword = flashlist.find("a", class_ = "SetPageTerm-wordText").text
        flashdefinition = flashlist.find("a", class_ = "SetPageTerm-definitionText").text
        with open(f"flashdata/{flashcard_title}.txt", "a") as file:
            file.write(f"w: {flashword}\n")
            file.write(f"def: {flashdefinition}\n")
    file.close()

root = Tk()

root.title("Flashcard_Studio/Scrape")
root.minsize(700,100)
root.maxsize(700,100)

label = Label(text = "Quizlet Link:")
entry_text = StringVar()
entry = Entry(textvariable = entry_text, width = 75)
button1 = Button(text = "Get Flashcards", command = scrape)
button2 = Button(root, text = "Review Sets", command = review_button)
label.pack()
entry.pack()
button1.pack()
button2.pack()

root.mainloop()


import pandas as pd
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
from sklearn.metrics import confusion_matrix
from sklearn.metrics import accuracy_score

# Charger les données dans un DataFrame Pandas
df = pd.read_csv('data/LDA_transformed_data.csv')

#Neural network
couches = 80
from sklearn.neural_network import MLPClassifier
X = df.iloc[:,1:].values
y = df.iloc[:,0].values.astype(int)
clf = MLPClassifier(solver='lbfgs', alpha=1e-5, hidden_layer_sizes=(couches,), random_state=42, max_iter=10000)
clf.fit(X, y)
MLPClassifier(alpha=1e-05, hidden_layer_sizes=(couches,), random_state=42, solver='lbfgs')

#Import well_data
df_well = pd.read_csv('data/well_data_with_facies.csv')

#LDA of well_data
# Sélectionner les variables d'entrée et la variable cible
X = df_well[['GR', 'ILD_log10', 'DeltaPHI', 'PHIND', 'PE']]
y = df_well['Facies']

# Transformation des données avec la LDA du training data
data = pd.read_csv('data/training_data.csv')
databis=data.iloc[:,0]

# Sélectionner les variables d'entrée et la variable cible
X_training = data[['GR', 'ILD_log10', 'DeltaPHI', 'PHIND', 'PE']]
y_training = data['Facies']

# Effectuer une AFD
lda = LinearDiscriminantAnalysis()
n=lda.fit(X_training, y_training)

#Transformer les données à prédire avec le même LDA que celui du training
ldafact_well=lda.transform(X)

#Prediction des données
y_pred = clf.predict(ldafact_well)

#print(confusion_matrix(y, y_pred))
print("LDA accuracy: " + str(accuracy_score(y, y_pred)))


"""
code by Tae Hwan Jung(Jeff Jung) @graykode, Derek Miller @dmmiller612, modify by shwei
Reference: https://github.com/jadore801120/attention-is-all-you-need-pytorch
           https://github.com/JayParks/transformer
"""

"""
代码位置：
手把手教你用Pytorch代码实现Transformer模型（超详细的代码解读）
https://ugirc.blog.csdn.net/article/details/120394042

"""

# ====================================================================================================
# 数据构建
import math
import torch
import numpy as np
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as Data

device = 'cpu'
# device = 'cuda'

# transformer epochs
epochs = 100
# epochs = 1000

# 这里我没有用什么大型的数据集，而是手动输入了两对德语→英语的句子
# 还有每个字的索引也是我手动硬编码上去的，主要是为了降低代码阅读难度
# S: Symbol that shows starting of decoding input
# E: Symbol that shows starting of decoding output
# P: Symbol that will fill in blank sequence if current batch data size is short than time steps
sentences = [
    # 德语和英语的单词个数不要求相同
    # enc_input                dec_input           dec_output
    ['ich mochte ein bier P', 'S i want a beer .', 'i want a beer . E'],
    ['ich mochte ein cola P', 'S i want a coke .', 'i want a coke . E']
]

# 德语和英语的单词要分开建立词库
# Padding Should be Zero
src_vocab = {'P': 0, 'ich': 1, 'mochte': 2, 'ein': 3, 'bier': 4, 'cola': 5}
src_idx2word = {i: w for i, w in enumerate(src_vocab)}
src_vocab_size = len(src_vocab)

tgt_vocab = {'P': 0, 'i': 1, 'want': 2, 'a': 3, 'beer': 4, 'coke': 5, 'S': 6, 'E': 7, '.': 8}
idx2word = {i: w for i, w in enumerate(tgt_vocab)}
tgt_vocab_size = len(tgt_vocab)

src_len = 5  # （源句子的长度）enc_input max sequence length
tgt_len = 6  # dec_input(=dec_output) max sequence length

# Transformer Parameters
d_model = 512  # Embedding Size（token embedding和position编码的维度）
d_ff = 2048  # FeedForward dimension (两次线性层中的隐藏层 512->2048->512，线性层是用来做特征提取的），当然最后会再接一个projection层
d_k = d_v = 64  # dimension of K(=Q), V（Q和K的维度需要相同，这里为了方便让K=V）
n_layers = 6  # number of Encoder of Decoder Layer（Block的个数）
n_heads = 8  # number of heads in Multi-Head Attention（有几套头）

# ==============================================================================================
# 数据构建

def make_data(sentences):
    """把单词序列转换为数字序列"""
    enc_inputs, dec_inputs, dec_outputs = [], [], []
    for i in range(len(sentences)):
        enc_input = [[src_vocab[n] for n in sentences[i][0].split()]]  # [[1, 2, 3, 4, 0], [1, 2, 3, 5, 0]]
        dec_input = [[tgt_vocab[n] for n in sentences[i][1].split()]]  # [[6, 1, 2, 3, 4, 8], [6, 1, 2, 3, 5, 8]]
        dec_output = [[tgt_vocab[n] for n in sentences[i][2].split()]]  # [[1, 2, 3, 4, 8, 7], [1, 2, 3, 5, 8, 7]]

        enc_inputs.extend(enc_input)
        dec_inputs.extend(dec_input)
        dec_outputs.extend(dec_output)

    return torch.LongTensor(enc_inputs), torch.LongTensor(dec_inputs), torch.LongTensor(dec_outputs)

enc_inputs, dec_inputs, dec_outputs = make_data(sentences)

class MyDataSet(Data.Dataset):
    """自定义DataLoader"""

    def __init__(self, enc_inputs, dec_inputs, dec_outputs):
        super(MyDataSet, self).__init__()
        self.enc_inputs = enc_inputs
        self.dec_inputs = dec_inputs
        self.dec_outputs = dec_outputs

    def __len__(self):
        return self.enc_inputs.shape[0]

    def __getitem__(self, idx):
        return self.enc_inputs[idx], self.dec_inputs[idx], self.dec_outputs[idx]

loader = Data.DataLoader(MyDataSet(enc_inputs, dec_inputs, dec_outputs), 2, True)

# ====================================================================================================
# Transformer模型

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        """
        x: [seq_len, batch_size, d_model]
        """
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)

def get_attn_pad_mask(seq_q, seq_k):
    # pad mask的作用：在对value向量加权平均的时候，可以让pad对应的alpha_ij=0，这样注意力就不会考虑到pad向量
    """这里的q,k表示的是两个序列（跟注意力机制的q,k没有关系），例如encoder_inputs (x1,x2,..xm)和encoder_inputs (x1,x2..xm)
    encoder和decoder都可能调用这个函数，所以seq_len视情况而定
    seq_q: [batch_size, seq_len]
    seq_k: [batch_size, seq_len]
    seq_len could be src_len or it could be tgt_len
    seq_len in seq_q and seq_len in seq_k maybe not equal
    """
    batch_size, len_q = seq_q.size()  # 这个seq_q只是用来expand维度的
    batch_size, len_k = seq_k.size()
    # eq(zero) is PAD token
    # 例如:seq_k = [[1,2,3,4,0], [1,2,3,5,0]]
    pad_attn_mask = seq_k.data.eq(0).unsqueeze(1)  # [batch_size, 1, len_k], True is masked
    return pad_attn_mask.expand(batch_size, len_q, len_k)  # [batch_size, len_q, len_k] 构成一个立方体(batch_size个这样的矩阵)

def get_attn_subsequence_mask(seq):
    """建议打印出来看看是什么的输出（一目了然）
    seq: [batch_size, tgt_len]
    """
    attn_shape = [seq.size(0), seq.size(1), seq.size(1)]
    # attn_shape: [batch_size, tgt_len, tgt_len]
    subsequence_mask = np.triu(np.ones(attn_shape), k=1)  # 生成一个上三角矩阵
    subsequence_mask = torch.from_numpy(subsequence_mask).byte()
    return subsequence_mask  # [batch_size, tgt_len, tgt_len]

# ==========================================================================================
class ScaledDotProductAttention(nn.Module):
    def __init__(self):
        super(ScaledDotProductAttention, self).__init__()

    def forward(self, Q, K, V, attn_mask):
        """
        Q: [batch_size, n_heads, len_q, d_k]
        K: [batch_size, n_heads, len_k, d_k]
        V: [batch_size, n_heads, len_v(=len_k), d_v]
        attn_mask: [batch_size, n_heads, seq_len, seq_len]
        说明：在encoder-decoder的Attention层中len_q(q1,..qt)和len_k(k1,...km)可能不同
        """
        scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(d_k)  # scores : [batch_size, n_heads, len_q, len_k]
        # mask矩阵填充scores（用-1e9填充scores中与attn_mask中值为1位置相对应的元素）
        scores.masked_fill_(attn_mask, -1e9)  # Fills elements of self tensor with value where mask is True.

        attn = nn.Softmax(dim=-1)(scores)  # 对最后一个维度(v)做softmax
        # scores : [batch_size, n_heads, len_q, len_k] * V: [batch_size, n_heads, len_v(=len_k), d_v]
        context = torch.matmul(attn, V)  # context: [batch_size, n_heads, len_q, d_v]
        # context：[[z1,z2,...],[...]]向量, attn注意力稀疏矩阵（用于可视化的）
        return context, attn

class MultiHeadAttention(nn.Module):
    """这个Attention类可以实现:
    Encoder的Self-Attention
    Decoder的Masked Self-Attention
    Encoder-Decoder的Attention
    """

    def __init__(self):
        super(MultiHeadAttention, self).__init__()
        self.W_Q = nn.Linear(d_model, d_k * n_heads, bias=False)  # q,k必须维度相同，不然无法做点积
        self.W_K = nn.Linear(d_model, d_k * n_heads, bias=False)
        self.W_V = nn.Linear(d_model, d_v * n_heads, bias=False)
        self.fc = nn.Linear(n_heads * d_v, d_model, bias=False)

    def forward(self, input_Q, input_K, input_V, attn_mask):
        """
        input_Q: [batch_size, len_q, d_model]
        input_K: [batch_size, len_k, d_model]
        input_V: [batch_size, len_v(=len_k), d_model]
        attn_mask: [batch_size, seq_len, seq_len]
        """
        residual, batch_size = input_Q, input_Q.size(0)
        # 下面的多头的参数矩阵是放在一起做线性变换的，然后再拆成多个头，这是工程实现的技巧
        # B: batch_size, S:seq_len, D: dim
        # (B, S, D) -proj-> (B, S, D_new) -split-> (B, S, Head, W) -trans-> (B, Head, S, W)
        #           线性变换               拆成多头

        # Q: [batch_size, n_heads, len_q, d_k]
        Q = self.W_Q(input_Q).view(batch_size, -1, n_heads, d_k).transpose(1, 2)
        # K: [batch_size, n_heads, len_k, d_k] # K和V的长度一定相同，维度可以不同
        K = self.W_K(input_K).view(batch_size, -1, n_heads, d_k).transpose(1, 2)
        # V: [batch_size, n_heads, len_v(=len_k), d_v]
        V = self.W_V(input_V).view(batch_size, -1, n_heads, d_v).transpose(1, 2)

        # 因为是多头，所以mask矩阵要扩充成4维的
        # attn_mask: [batch_size, seq_len, seq_len] -> [batch_size, n_heads, seq_len, seq_len]
        attn_mask = attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1)

        # context: [batch_size, n_heads, len_q, d_v], attn: [batch_size, n_heads, len_q, len_k]
        context, attn = ScaledDotProductAttention()(Q, K, V, attn_mask)
        # 下面将不同头的输出向量拼接在一起
        # context: [batch_size, n_heads, len_q, d_v] -> [batch_size, len_q, n_heads * d_v]
        context = context.transpose(1, 2).reshape(batch_size, -1, n_heads * d_v)
        # 再做一个projection
        output = self.fc(context)  # [batch_size, len_q, d_model]
        return nn.LayerNorm(d_model).to(device)(output + residual), attn

# Pytorch中的Linear只会对最后一维操作，所以正好是我们希望的每个位置用同一个全连接网络
class PoswiseFeedForwardNet(nn.Module):
    def __init__(self):
        super(PoswiseFeedForwardNet, self).__init__()
        self.fc = nn.Sequential(
            nn.Linear(d_model, d_ff, bias=False),
            nn.ReLU(),
            nn.Linear(d_ff, d_model, bias=False)
        )

    def forward(self, inputs):
        """
        inputs: [batch_size, seq_len, d_model]
        """
        residual = inputs
        output = self.fc(inputs)
        return nn.LayerNorm(d_model).to(device)(output + residual)  # [batch_size, seq_len, d_model]

class EncoderLayer(nn.Module):
    def __init__(self):
        super(EncoderLayer, self).__init__()
        self.enc_self_attn = MultiHeadAttention()
        self.pos_ffn = PoswiseFeedForwardNet()

    def forward(self, enc_inputs, enc_self_attn_mask):
        """E
        enc_inputs: [batch_size, src_len, d_model]
        enc_self_attn_mask: [batch_size, src_len, src_len]  mask矩阵(pad mask or sequence mask)
        """
        # enc_outputs: [batch_size, src_len, d_model], attn: [batch_size, n_heads, src_len, src_len]
        # 第一个enc_inputs * W_Q = Q
        # 第二个enc_inputs * W_K = K
        # 第三个enc_inputs * W_V = V
        enc_outputs, attn = self.enc_self_attn(enc_inputs, enc_inputs, enc_inputs,
                                               enc_self_attn_mask)  # enc_inputs to same Q,K,V（未线性变换前）
        enc_outputs = self.pos_ffn(enc_outputs)
        # enc_outputs: [batch_size, src_len, d_model]
        return enc_outputs, attn

class DecoderLayer(nn.Module):
    def __init__(self):
        super(DecoderLayer, self).__init__()
        self.dec_self_attn = MultiHeadAttention()
        self.dec_enc_attn = MultiHeadAttention()
        self.pos_ffn = PoswiseFeedForwardNet()

    def forward(self, dec_inputs, enc_outputs, dec_self_attn_mask, dec_enc_attn_mask):
        """
        dec_inputs: [batch_size, tgt_len, d_model]
        enc_outputs: [batch_size, src_len, d_model]
        dec_self_attn_mask: [batch_size, tgt_len, tgt_len]
        dec_enc_attn_mask: [batch_size, tgt_len, src_len]
        """
        # dec_outputs: [batch_size, tgt_len, d_model], dec_self_attn: [batch_size, n_heads, tgt_len, tgt_len]
        dec_outputs, dec_self_attn = self.dec_self_attn(dec_inputs, dec_inputs, dec_inputs,
                                                        dec_self_attn_mask)  # 这里的Q,K,V全是Decoder自己的输入
        # dec_outputs: [batch_size, tgt_len, d_model], dec_enc_attn: [batch_size, h_heads, tgt_len, src_len]
        dec_outputs, dec_enc_attn = self.dec_enc_attn(dec_outputs, enc_outputs, enc_outputs,
                                                      dec_enc_attn_mask)  # Attention层的Q(来自decoder) 和 K,V(来自encoder)
        dec_outputs = self.pos_ffn(dec_outputs)  # [batch_size, tgt_len, d_model]
        return dec_outputs, dec_self_attn, dec_enc_attn  # dec_self_attn, dec_enc_attn这两个是为了可视化的

class Encoder(nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()
        self.src_emb = nn.Embedding(src_vocab_size, d_model)  # token Embedding
        self.pos_emb = PositionalEncoding(d_model)  # Transformer中位置编码时固定的，不需要学习
        self.layers = nn.ModuleList([EncoderLayer() for _ in range(n_layers)])

    def forward(self, enc_inputs):
        """
        enc_inputs: [batch_size, src_len]
        """
        enc_outputs = self.src_emb(enc_inputs)  # [batch_size, src_len, d_model]
        enc_outputs = self.pos_emb(enc_outputs.transpose(0, 1)).transpose(0, 1)  # [batch_size, src_len, d_model]
        # Encoder输入序列的pad mask矩阵
        enc_self_attn_mask = get_attn_pad_mask(enc_inputs, enc_inputs)  # [batch_size, src_len, src_len]
        enc_self_attns = []  # 在计算中不需要用到，它主要用来保存你接下来返回的attention的值（这个主要是为了你画热力图等，用来看各个词之间的关系
        for layer in self.layers:  # for循环访问nn.ModuleList对象
            # 上一个block的输出enc_outputs作为当前block的输入
            # enc_outputs: [batch_size, src_len, d_model], enc_self_attn: [batch_size, n_heads, src_len, src_len]
            enc_outputs, enc_self_attn = layer(enc_outputs,
                                               enc_self_attn_mask)  # 传入的enc_outputs其实是input，传入mask矩阵是因为你要做self attention
            enc_self_attns.append(enc_self_attn)  # 这个只是为了可视化
        return enc_outputs, enc_self_attns

class Decoder(nn.Module):
    def __init__(self):
        super(Decoder, self).__init__()
        self.tgt_emb = nn.Embedding(tgt_vocab_size, d_model)  # Decoder输入的embed词表
        self.pos_emb = PositionalEncoding(d_model)
        self.layers = nn.ModuleList([DecoderLayer() for _ in range(n_layers)])  # Decoder的blocks

    def forward(self, dec_inputs, enc_inputs, enc_outputs):
        """
        dec_inputs: [batch_size, tgt_len]
        enc_inputs: [batch_size, src_len]
        enc_outputs: [batch_size, src_len, d_model]   # 用在Encoder-Decoder Attention层
        """
        dec_outputs = self.tgt_emb(dec_inputs)  # [batch_size, tgt_len, d_model]
        dec_outputs = self.pos_emb(dec_outputs.transpose(0, 1)).transpose(0, 1).to(
            device)  # [batch_size, tgt_len, d_model]
        # Decoder输入序列的pad mask矩阵（这个例子中decoder是没有加pad的，实际应用中都是有pad填充的）
        dec_self_attn_pad_mask = get_attn_pad_mask(dec_inputs, dec_inputs).to(device)  # [batch_size, tgt_len, tgt_len]
        # Masked Self_Attention：当前时刻是看不到未来的信息的
        dec_self_attn_subsequence_mask = get_attn_subsequence_mask(dec_inputs).to(
            device)  # [batch_size, tgt_len, tgt_len]

        # Decoder中把两种mask矩阵相加（既屏蔽了pad的信息，也屏蔽了未来时刻的信息）
        dec_self_attn_mask = torch.gt((dec_self_attn_pad_mask + dec_self_attn_subsequence_mask),
                                      0).to(device)  # [batch_size, tgt_len, tgt_len]; torch.gt比较两个矩阵的元素，大于则返回1，否则返回0

        # 这个mask主要用于encoder-decoder attention层
        # get_attn_pad_mask主要是enc_inputs的pad mask矩阵(因为enc是处理K,V的，求Attention时是用v1,v2,..vm去加权的，要把pad对应的v_i的相关系数设为0，这样注意力就不会关注pad向量)
        #                       dec_inputs只是提供expand的size的
        dec_enc_attn_mask = get_attn_pad_mask(dec_inputs, enc_inputs)  # [batc_size, tgt_len, src_len]

        dec_self_attns, dec_enc_attns = [], []
        for layer in self.layers:
            # dec_outputs: [batch_size, tgt_len, d_model], dec_self_attn: [batch_size, n_heads, tgt_len, tgt_len], dec_enc_attn: [batch_size, h_heads, tgt_len, src_len]
            # Decoder的Block是上一个Block的输出dec_outputs（变化）和Encoder网络的输出enc_outputs（固定）
            dec_outputs, dec_self_attn, dec_enc_attn = layer(dec_outputs, enc_outputs, dec_self_attn_mask,
                                                             dec_enc_attn_mask)
            dec_self_attns.append(dec_self_attn)
            dec_enc_attns.append(dec_enc_attn)
        # dec_outputs: [batch_size, tgt_len, d_model]
        return dec_outputs, dec_self_attns, dec_enc_attns

class Transformer(nn.Module):
    def __init__(self):
        super(Transformer, self).__init__()
        self.encoder = Encoder().to(device)
        self.decoder = Decoder().to(device)
        self.projection = nn.Linear(d_model, tgt_vocab_size, bias=False).to(device)

    def forward(self, enc_inputs, dec_inputs):
        """Transformers的输入：两个序列
        enc_inputs: [batch_size, src_len]
        dec_inputs: [batch_size, tgt_len]
        """
        # tensor to store decoder outputs
        # outputs = torch.zeros(batch_size, tgt_len, tgt_vocab_size).to(self.device)

        # enc_outputs: [batch_size, src_len, d_model], enc_self_attns: [n_layers, batch_size, n_heads, src_len, src_len]
        # 经过Encoder网络后，得到的输出还是[batch_size, src_len, d_model]
        enc_outputs, enc_self_attns = self.encoder(enc_inputs)
        # dec_outputs: [batch_size, tgt_len, d_model], dec_self_attns: [n_layers, batch_size, n_heads, tgt_len, tgt_len], dec_enc_attn: [n_layers, batch_size, tgt_len, src_len]
        dec_outputs, dec_self_attns, dec_enc_attns = self.decoder(dec_inputs, enc_inputs, enc_outputs)
        # dec_outputs: [batch_size, tgt_len, d_model] -> dec_logits: [batch_size, tgt_len, tgt_vocab_size]
        dec_logits = self.projection(dec_outputs)
        return dec_logits.view(-1, dec_logits.size(-1)), enc_self_attns, dec_self_attns, dec_enc_attns

model = Transformer().to(device)
# 这里的损失函数里面设置了一个参数 ignore_index=0，因为 "pad" 这个单词的索引为 0，这样设置以后，就不会计算 "pad" 的损失（因为本来 "pad" 也没有意义，不需要计算）
criterion = nn.CrossEntropyLoss(ignore_index=0)
optimizer = optim.SGD(model.parameters(), lr=1e-3, momentum=0.99)  # 用adam的话效果不好

# ====================================================================================================
for epoch in range(epochs):
    for enc_inputs, dec_inputs, dec_outputs in loader:
        """
        enc_inputs: [batch_size, src_len]
        dec_inputs: [batch_size, tgt_len]
        dec_outputs: [batch_size, tgt_len]
        """
        enc_inputs, dec_inputs, dec_outputs = enc_inputs.to(device), dec_inputs.to(device), dec_outputs.to(device)
        # outputs: [batch_size * tgt_len, tgt_vocab_size]
        outputs, enc_self_attns, dec_self_attns, dec_enc_attns = model(enc_inputs, dec_inputs)
        loss = criterion(outputs, dec_outputs.view(-1))  # dec_outputs.view(-1):[batch_size * tgt_len * tgt_vocab_size]
        print('Epoch:', '%04d' % (epoch + 1), 'loss =', '{:.6f}'.format(loss))

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

def greedy_decoder(model, enc_input, start_symbol):
    """贪心编码
    For simplicity, a Greedy Decoder is Beam search when K=1. This is necessary for inference as we don't know the
    target sequence input. Therefore we try to generate the target input word by word, then feed it into the transformer.
    Starting Reference: http://nlp.seas.harvard.edu/2018/04/03/attention.html#greedy-decoding
    :param model: Transformer Model
    :param enc_input: The encoder input
    :param start_symbol: The start symbol. In this example it is 'S' which corresponds to index 4
    :return: The target input
    """
    enc_outputs, enc_self_attns = model.encoder(enc_input)
    dec_input = torch.zeros(1, 0).type_as(enc_input.data)
    terminal = False
    next_symbol = start_symbol
    while not terminal:
        # 预测阶段：dec_input序列会一点点变长（每次添加一个新预测出来的单词）
        dec_input = torch.cat([dec_input.to(device), torch.tensor([[next_symbol]], dtype=enc_input.dtype).to(device)],
                              -1)
        dec_outputs, _, _ = model.decoder(dec_input, enc_input, enc_outputs)
        projected = model.projection(dec_outputs)
        prob = projected.squeeze(0).max(dim=-1, keepdim=False)[1]
        # 增量更新（我们希望重复单词预测结果是一样的）
        # 我们在预测是会选择性忽略重复的预测的词，只摘取最新预测的单词拼接到输入序列中
        next_word = prob.data[-1]  # 拿出当前预测的单词(数字)。我们用x'_t对应的输出z_t去预测下一个单词的概率，不用z_1,z_2..z_{t-1}
        next_symbol = next_word
        if next_symbol == tgt_vocab["E"]:
            terminal = True
        # print(next_word)

    # greedy_dec_predict = torch.cat(
    #     [dec_input.to(device), torch.tensor([[next_symbol]], dtype=enc_input.dtype).to(device)],
    #     -1)
    greedy_dec_predict = dec_input[:, 1:]
    return greedy_dec_predict

# ==========================================================================================
# 预测阶段
enc_inputs, _, _ = next(iter(loader))
for i in range(len(enc_inputs)):
    greedy_dec_predict = greedy_decoder(model, enc_inputs[i].view(1, -1).to(device), start_symbol=tgt_vocab["S"])
    print(enc_inputs[i], '->', greedy_dec_predict.squeeze())
    print([src_idx2word[t.item()] for t in enc_inputs[i]], '->',
          [idx2word[n.item()] for n in greedy_dec_predict.squeeze()])


import os
from datetime import datetime
from pprint import pprint as pp
import numpy as np
import tensorflow as tf
from sklearn.model_selection import train_test_split
from common import *
from optparse import OptionParser

# LEARNING_RATE = 0.02
LEARNING_RATE = 0.1
DATASET_SIZE = 1000
EPOCHS = 1200

# Collection of dataset
n = 0.25

X = np.linspace(0, 0.5 , DATASET_SIZE)

y = np.array(
    list([np.sin(n * x * np.pi)] for x in X)
)

x_train, x_test, y_train, y_test = train_test_split(X, y, test_size = 0.2)
activation_function = 'tanh'

model = tf.keras.Sequential([
        tf.keras.layers.Dense(10, input_dim=1, activation=activation_function),
        tf.keras.layers.Dense(5, activation=activation_function),
        tf.keras.layers.Dense(1, activation=activation_function)
    ])

optimizer = tf.keras.optimizers.SGD(learning_rate=LEARNING_RATE)

model.compile(optimizer=optimizer, loss="MSE")

#Training
model.fit(x_train, y_train, epochs=EPOCHS)

save = True

if save:
    timestamp = datetime.now()
    timestamp = timestamp.strftime("%d-%m-%Y--%H-%M-%S-%f")

    dir_path = f"./saved_models/ANN3/{timestamp}"

    os.mkdir(dir_path)

    model.save(f"{dir_path}/model")

    with open(f'{dir_path}/metadata.txt', 'w') as f:
        f.write(f"LEARNING_RATE = {LEARNING_RATE}, DATASET_SIZE = {DATASET_SIZE}, EPOCHS = {EPOCHS}\n")
        f.write(f"ACTIVATION FUNCTION: '{activation_function}'\n")
        f.write(str(model.evaluate(x_test, y_test)))
        f.write('\n')


import os
import time
import sys
import traceback
import glob

import cv2
import numpy as np
import onnxruntime

def main(args):
    classes = ['nsfw', 'sfw']

    session =  onnxruntime.InferenceSession(args.model_path, providers=['CUDAExecutionProvider'])
    input_name = session.get_inputs()[0].name
    output_name = [output.name for output in session.get_outputs()]

    tp, fp, tn, fn = 0, 0, 0, 0
    total_num = 0
    total_used_time = 0.0
    for ipath in glob.glob("/data/test/*/*.jpg", recursive=True):
        image = cv2.imread(ipath)
        if image is not None:
            total_num += 1
            label = ipath.split('/')[-2]
            if total_num >= 5:
                bg_time = time.time()
            blob = cv2.dnn.blobFromImage(image, 1.0/127.5, (224, 224), (127.5, 127.5, 127.5), swapRB=True)
            out = session.run(output_name, {input_name: blob})[0][0]
            if total_num >= 5:
                total_used_time += (time.time() - bg_time)
            pred = classes[np.argmax(out)]
            if pred == label == "positive":
                tp += 1
            elif pred == label == "negative":
                tn += 1
            elif pred == "negative" and label == "positive":
                fn += 1
            else:
                fp += 1
    print('Total test number is: {}, average used time is: {}'.format(total_num, total_used_time/(total_num-4)))
    """
    FPR: False Positive Rate,  FPR = FP / N = FP / (FP+TN) = 1 - TNR
    FNR: False Negtive Rate, FNR = FN / P = FN / (FN+TP) = 1 - TPR
    TPR: True Positive Rate, TPR = TP / p = TP / (TP+FN) = 1 - FNR
    TNR: True Negtive Rate, TNR = TN / N = TN / (TN+FP) = 1 - FPR
    Accuracy = (TP+TN) / (TP+FP+FN+TN)
    Precision(PPV) = TP / (TP+FP)
    Recall(TPR) = TP / (TP+FN) = 1 - FNR
    F1 score = 2 * precision * recall / (precision + recall)
    要求FPR和FNR越低越好
    """

    fpr = fp / (fp + tn)
    fnr = fn / (fn + tp)
    accuracy = (tp + tn) / (tp + fp + fn + tn)
    precision = tp / (tp + fp)
    recall = tp / (tp + fn)
    f1score = 2 * precision * recall / (precision + recall)

    print('FPR={}, FNR={}, accuracy={}, precision={}, recall={}, f1score={}'.format
          (fpr, fnr, accuracy, precision, recall, f1score))

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description='ONNX GPU Inference')

    parser.add_argument('--model-path', default='./models/model.onnx', help='onnx model')
    parser.add_argument('--resize', default=224, type=int, help='size of resize')

    args = parser.parse_args()

    print(args)
    main(args)


import ase.io
import fire

from ai2_kit.core.log import get_logger
from ai2_kit import res
from ai2_kit.domain.cp2k import dump_coord_n_cell
from typing import Optional, Literal
from ase import Atoms, Atom
from string import Template
import re
import os
import json

logger = get_logger(__name__)

DEEPMD_DEFAULT_TEMPLATE = os.path.join(res.DIR_PATH, 'catalysis/deepmd.json')
MLP_TRAINING_TEMPLATE = os.path.join(res.DIR_PATH, 'catalysis/mlp-training.yml')

CP2K_DEFAULT_TEMPLATE = os.path.join(res.DIR_PATH, 'catalysis/cp2k.inp')

CP2K_SCF_SEMICONDUCTOR = """\
        # CONFIGURATION FOR SEMICONDUCTOR
        &SCF
            SCF_GUESS RESTART
            EPS_SCF 3e-07
            MAX_SCF 50
            &OUTER_SCF
                EPS_SCF 3e-07
                MAX_SCF 20
            &END OUTER_SCF
            &OT
                MINIMIZER CG
                PRECONDITIONER FULL_SINGLE_INVERSE
                ENERGY_GAP 0.1
            &END OT
        &END SCF
        # END CONFIGURATION FOR SEMICONDUCTOR"""

CP2K_SCF_METAL = """\
        # CONFIGURATION FOR METAL
        &SCF
            SCF_GUESS RESTART
            EPS_SCF 3e-07
            MAX_SCF 500
            ADDED_MOS 500
            CHOLESKY INVERSE
            &SMEAR
                METHOD FERMI_DIRAC
                ELECTRONIC_TEMPERATURE [K] 300
            &END SMEAR
            &DIAGONALIZATION
                ALGORITHM STANDARD
            &END DIAGONALIZATION
            &MIXING
                METHOD BROYDEN_MIXING
                ALPHA 0.3
                BETA 1.5
                NBROYDEN 14
            &END MIXING
        &END SCF
        # END CONFIGURATION FOR METAL"""

CP2K_MOTION_TEMPLATE = """\
&MOTION
  &MD
    ENSEMBLE NVT
    STEPS       $steps
    TIMESTEP    $timestep
    TEMPERATURE $temp
    &THERMOSTAT
       TYPE CSVR
       REGION MASSIVE
       &CSVR
          TIMECON [fs] 100.0
       &END
    &END
  &END MD
  &PRINT
   &TRAJECTORY
     &EACH
       MD 1
     &END EACH
   &END TRAJECTORY
   &VELOCITIES
     &EACH
       MD 1
     &END EACH
   &END VELOCITIES
   &FORCES
     &EACH
       MD 1
     &END EACH
   &END FORCES
   &RESTART_HISTORY
     &EACH
       MD 1000
     &END EACH
   &END RESTART_HISTORY
   &RESTART
     BACKUP_COPIES 3
     &EACH
       MD 1
     &END EACH
   &END RESTART
  &END PRINT
&END MOTION"""

CP2K_SCF_TABLE = {
    'metal': CP2K_SCF_METAL,
    'semi': CP2K_SCF_SEMICONDUCTOR,
}

CP2K_ACCURACY_TABLE = {
    'high': {'cutoff': 1000, 'rel_cutoff': 90 },
    'medium': {'cutoff': 800, 'rel_cutoff': 70 },
    'low': {'cutoff': 600, 'rel_cutoff': 50 },
}

class ConfigBuilder:

    def __init__(self):
        self._atoms: Optional[Atoms] = None

    def load_system(self, file: str, **kwargs):
        """
        Loading system to memory using ASE
        """
        assert self._atoms is None, 'atoms is already loaded'
        self._atoms = ase.io.read(file, **kwargs)  # type: ignore
        return self

    def gen_mlp_training_input(self,
                               out_dir: str = 'out',
                               template_file: str = MLP_TRAINING_TEMPLATE):
        # Read yaml file as text so that the comments are preserved
        with open(template_file, 'r') as fp:
            text = fp.read()
        # Generate the type_map and mass_map automatically
        assert self._atoms is not None, 'atoms must be loaded first'
        type_map = list(set(self._atoms.get_chemical_symbols()))
        mass_map = [Atom(symbol).mass for symbol in type_map]

        out_data = Template(text).substitute(
            type_map=json.dumps(type_map),
            mass_map=json.dumps(mass_map),
        )
        os.makedirs(out_dir, exist_ok=True)
        mlp_training_input_path = os.path.join(out_dir, 'mlp-training.yml')
        with open(mlp_training_input_path, 'w', encoding='utf-8') as fp:
            fp.write(out_data)

    def gen_deepmd_input(self,
                         out_dir: str = 'out',
                         steps: int = 10000,
                         template_file: str = DEEPMD_DEFAULT_TEMPLATE):
        with open(template_file, 'r') as fp:
            data = json.load(fp)
        os.makedirs(out_dir, exist_ok=True)

        data['training']['numb_steps'] = steps
        deepmd_input_path = os.path.join(out_dir, 'deepmd.json')
        with open(deepmd_input_path, 'w', encoding='utf-8') as fp:
            json.dump(data, fp, indent=4)

    def gen_plumed_input(self, out_dir: str = 'out'):
        """
        Generate Plumed input files
        Args:
            out_dir: output directory, plumed.dat will be generated in this directory
        """
        assert self._atoms is not None, 'atoms must be loaded first'
        plumed_input = [ 'UNITS LENGTH=A' ]
        # define atoms groups by element, e.g:
        # Ag: GROUP ATOMS=1,2,3,4,5,6,7,8,9,10
        # O: GROUP ATOMS=11,12,13,14,15,16,17,18,19,20
        element2ids = {}
        for i, element in enumerate(self._atoms.get_chemical_symbols(), start=1):
            element2ids.setdefault(element, []).append(i)
        for element, ids in element2ids.items():
            plumed_input.append(f'{element}: GROUP ATOMS={",".join(map(str, ids))}')

        # define reaction coordinate
        plumed_input.extend([
            '#define more groups if needed',
            '',
            '# define reaction coordinates, e.g. CV1, CV2, ...',
            '# you may define as many as you want',
            'CV1:',
            '',
            '# define sampling method: metadynamics',
            'metad: METAD ARG=CV1 SIGMA=0.1 HEIGHT=5 PACE=100 TEMP=1000 FILE=HILLS',
            '# define more commands if you need',
            '',
            '# print CVs',
            'PRINT STRIDE=10 ARG=CV1,metad.bias FILE=COLVAR',
        ])
        os.makedirs(out_dir, exist_ok=True)
        plumed_input_path = os.path.join(out_dir, 'plumed.dat')
        with open(plumed_input_path, 'w', encoding='utf-8') as fp:
            fp.write('\n'.join(plumed_input))

    def gen_lammps_input(self):
        ... # TODO

    def gen_report(self):
        ... # TODO

    def gen_cp2k_input(self,
                       basic_set_file: str = 'BASIS_MOLOPT',
                       potential_file: str = 'GTH_POTENTIALS',
                       out_dir: str = 'out',
                       style: Literal['metal', 'semi'] = 'metal',
                       accuracy: Literal['high', 'medium', 'low'] = 'medium',
                       aimd: bool = False,
                       temp: float = 330.,
                       steps: int = 1000,
                       timestep: float = 0.5,
                       parameter_file: str = 'dftd3.dat',
                       template_file: str = CP2K_DEFAULT_TEMPLATE,):
        """
        Generate CP2K input files

        You should call `load_system` first to load system into memory.
        And you may also need to ensure the basic set and potential files
        you want to use are available in CP2K_DATA_DIR,
        or else you have to specify the full path instead of their file name.

        Args:
            basic_set_file: basic set file, can be path or file in CP2K_DATA_DIR
            potential_file: potential file, can be path or file in CP2K_DATA_DIR
            out_dir: output directory, cp2k.inp and coord_n_cell.inc will be generated in this directory
            style: 'metal' or 'semi'
            accuracy: 'high', 'medium' or 'low'
            aimd: whether to run AIMD
            temp: temperature for AIMD
            steps: steps for AIMD
            timestep: timestep for AIMD
            parameter_file: parameter file, can be path or file in CP2K_DATA_DIR
            template_file: template file, no need to change in most cases
        """

        assert self._atoms is not None, 'atoms must be loaded first'
        # get available basic set and potential
        with open(find_cp2k_data_file(basic_set_file), 'r') as fp:
            basic_set_table = parse_cp2k_basic_set_potential(fp)
        with open(find_cp2k_data_file(potential_file), 'r') as fp:
            potential_table = parse_cp2k_basic_set_potential(fp)
        with open(template_file, 'r') as fp:
            template = fp.read()
        # build kinds
        elements = set(self._atoms.get_chemical_symbols())

        kinds = ''
        total_ve = 0  # Valence Electron

        for element in elements:
            basic_set_all = basic_set_table.get(element)
            basic_set = basic_set_all[-1]
            logger.info(f'BASIC_SET {basic_set} is selected for {element}')

            potential_all = potential_table.get(element)
            potential = next(p for p in potential_all if 'PBE' in p)
            logger.info(f'POTENTIAL {potential} is selected for {element}')
            kinds += '\n'.join([
                f'    &KIND {element}',
                f'        BASIS_SET {basic_set}',
                f'        # All available BASIS_SET:',
                f'        # {" ".join(basic_set_all)}',
                f'        POTENTIAL {potential}',
                f'        # All available POTENTIAL:',
                f'        # {" ".join(potential_all)}',
                f'    &END KIND',
                '', # This empty line is required
            ])
            # get ve from the last number of potential
            # for example, if the potential is GTH-OLYP-q6
            # then the ve is 6
            ve = int(re.match(r'.*?(\d+)$', potential).group(1))
            total_ve += ve * self._atoms.get_chemical_symbols().count(element)
        logger.info("total valence electron of the system: %d" % total_ve)

        motion = ''
        if aimd:
            motion = Template(CP2K_MOTION_TEMPLATE).substitute(
                steps=steps,
                timestep=timestep,
                temp=temp,
            )

        template_vars = {
            'run_type': 'MD' if aimd else 'ENERGY_FORCE',
            'basic_set_file': basic_set_file,
            'potential_file': potential_file,
            'parameter_file': parameter_file,
            'uks': 'T' if total_ve % 2 else 'F',
            'kinds': kinds,
            'motion': motion,
            'scf': CP2K_SCF_TABLE[style],
            **CP2K_ACCURACY_TABLE[accuracy],
        }

        cp2k_input = Template(template).substitute(**template_vars)
        os.makedirs(out_dir, exist_ok=True)
        cp2k_input_path = os.path.join(out_dir, 'cp2k.inp')
        with open(cp2k_input_path, 'w') as fp:
            fp.write(cp2k_input)
        logger.info(f'CP2K input file is generated at {cp2k_input_path}')

        coord_n_cell_path = os.path.join(out_dir, 'coord_n_cell.inc')
        with open(coord_n_cell_path, 'w') as fp:
            dump_coord_n_cell(fp, self._atoms)
        logger.info(f'coord_n_cell.inc is generated at {coord_n_cell_path}')

def find_cp2k_data_file(file: str):
    """
    find CP2K data file in CP2K_DATA_DIR
    """
    if not os.path.exists(file):
        cp2k_data_dir = os.environ.get('CP2K_DATA_DIR')
        if cp2k_data_dir is None:
            raise FileNotFoundError(f'Cannot find {file}')
        file = os.path.join(cp2k_data_dir, file)
        if not os.path.exists(file):
            raise FileNotFoundError(f'Cannot find {file}')
    return file

def parse_cp2k_basic_set_potential(fp):
    """
    get all available basic set and potential from CP2K data file
    """
    parsed = {}
    for line in fp:
        line: str = line.strip()
        if line.startswith('#'):
            continue # skip comment
        tokens = line.split()
        if 0 == len(tokens):
            continue  # skip empty line
        if re.match(r'^[A-Z][a-z]*$', tokens[0]):
            parsed.setdefault(tokens[0], []).append(tokens[1])
    return parsed

class CmdEntries:

    def build_config(self):
        """
        Build config file for catalyst
        """
        return ConfigBuilder

def cli_main():
    fire.Fire(CmdEntries)

if __name__ == '__main__':
    fire.Fire(ConfigBuilder)


"""Added double indexes

Revision ID: 5783809706d2
Revises: 4e2dde2d15e6
Create Date: 2015-06-01 21:08:15.106408

"""

# revision identifiers, used by Alembic.
revision = '5783809706d2'
down_revision = '4e2dde2d15e6'

from alembic import op
import sqlalchemy as sa

def upgrade():
    ### commands auto generated by Alembic - please adjust! ###
    op.create_index('ix_ecc_pi_ecc', 'ecc', ['pi', 'ecc'], unique=False)
    op.create_index('ix_station_spid_radioepg_enabled', 'station', ['service_provider_id', 'radioepg_enabled'], unique=False)
    ### end Alembic commands ###

def downgrade():
    ### commands auto generated by Alembic - please adjust! ###
    op.drop_index('ix_station_spid_radioepg_enabled', table_name='station')
    op.drop_index('ix_ecc_pi_ecc', table_name='ecc')
    ### end Alembic commands ###


import paddle
import itertools
import log
import sys
import time
import datetime
import paddle.vision.transforms as T
from paddle.vision.datasets import Cifar10
from paddle.optimizer import Adam
from paddle.optimizer.lr import CosineAnnealingDecay
from paddle.optimizer.lr import LinearWarmup
from paddle.regularizer import L2Decay
from paddle.nn import CrossEntropyLoss
from paddle.metric import Accuracy
from Edge_blur import EdgeBlurs
import ViT.module

class Cosine(CosineAnnealingDecay):
    def __init__(self, lr, step_each_epoch, epochs, **kwargs):
        super(Cosine, self).__init__(learning_rate=lr, T_max=step_each_epoch * epochs)
        self.update_specified = False

class CosineWarmup(LinearWarmup):
    def __init__(self, lr, step_each_epoch, epochs, warmup_epoch=5, **kwargs):
        assert epochs > warmup_epoch, "total epoch({}) should be larger than warmup_epoch({}) in CosineWarmup.".format(
            epochs, warmup_epoch)
        warmup_step = warmup_epoch * step_each_epoch
        start_lr = 0.0
        end_lr = lr
        lr_sch = Cosine(lr, step_each_epoch, epochs - warmup_epoch)

        super(CosineWarmup, self).__init__(learning_rate=lr_sch, warmup_steps=warmup_step, start_lr=start_lr,
                                           end_lr=end_lr)
        self.update_specified = False

sys.stdout = log.Logger('./log/log{}.txt'.format(
    time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time()))),
    mode='w', encoding='utf-8')

if __name__ == '__main__':
    start_time = datetime.datetime.now()
    print(f'训练开始，当前时间为{start_time}')

    print("=======begin=======")

    img_size = 128

    drop = 0.2

    net = ViT.module.VisionTransformer(
        img_size=img_size,
        patch_size=16,
        class_dim=10,
        embed_dim=384,
        depth=4,
        num_heads=6,
        drop_rate=drop,
        attn_drop_rate=drop,
        drop_path_rate=drop,
        mlp_ratio=4,
        qkv_bias=True,
        epsilon=1e-6)
    model = paddle.Model(net)
    model.load(r'pre param/63.pdparams')
    edge_blurs = EdgeBlurs()
    train_transform = T.Compose(
        [
            T.Resize(img_size),
            edge_blurs,
            T.RandomErasing(),
            T.RandomHorizontalFlip(),
            T.RandomVerticalFlip(),
            T.ToTensor(),
            T.Normalize(
                mean=[0.5, 0.5, 0.5],
                std=[0.5, 0.5, 0.5],
                to_rgb=True,
            ),
        ]
    )

    test_transform = T.Compose(
        [
            T.Resize(img_size),
            T.ToTensor(),
            T.Normalize(
                mean=[0.5, 0.5, 0.5],
                std=[0.5, 0.5, 0.5],
                to_rgb=True,
            ),
        ]
    )

    train_dataset = Cifar10(mode='train', transform=train_transform, download=True, backend='cv2')
    # print(len(train_dataset))
    test_dataset = Cifar10(mode='test', transform=test_transform, download=True, backend='cv2')
    # print(len(test_dataset))

    epochs = 300
    warmup_epoch = 10

    # for img, label in itertools.islice(iter(train_dataset), 5):
    #     print(type(img), img.shape, label)

    lr = 1e-3
    # scheduler = CosineWarmup(lr=lr, step_each_epoch=100, epochs=epochs, warmup_epoch=warmup_epoch, start_lr = 0, end_lr = lr, verbose = True)
    # scheduler = LinearWarmup(learning_rate=lr, warmup_steps=warmup_epoch, start_lr=0, end_lr=lr)
    # scheduler = CosineAnnealingDecay(learning_rate=lr, T_max=warmup_epoch, eta_min=0)
    optimizer = Adam(learning_rate=lr, parameters=model.parameters())
    model.prepare(optimizer, CrossEntropyLoss(), Accuracy(topk=(1, 5)))

    print()
    print(f'the lr is {lr}')
    print(f'the epochs is {epochs}, and the warmup epoch is {warmup_epoch}')
    print()

    use_gpu = True
    paddle.set_device('gpu:0')
    model.fit(train_dataset,
              test_dataset,
              epochs=epochs,
              batch_size=300,
              save_dir="param",
              save_freq=7,
              log_freq=1,
              num_workers=0,
              shuffle=True,
              eval_freq=7,
              verbose=1)

    print("========end========")

    end_time = datetime.datetime.now()
    print(f'训练结束，当前时间为{end_time}')


"""bitforum URL Configuration

The `urlpatterns` list routes URLs to views. For more information please see:
    https://docs.djangoproject.com/en/3.1/topics/http/urls/
Examples:
Function views
    1. Add an import:  from my_app import views
    2. Add a URL to urlpatterns:  path('', views.home, name='home')
Class-based views
    1. Add an import:  from other_app.views import Home
    2. Add a URL to urlpatterns:  path('', Home.as_view(), name='home')
Including another URLconf
    1. Import the include() function: from django.urls import include, path
    2. Add a URL to urlpatterns:  path('blog/', include('blog.urls'))
"""
from django.contrib import admin
from django.urls import path
from django.urls import path, include
from posts import urls
from .views import *
from django.conf import settings
from django.conf.urls.static import static

urlpatterns = [
    path('admin/', admin.site.urls),
    # path('admin',admin.site.urls),
    path('myProfile',myProfile),
    path('login',login_signup_page),
    path('checkLogin',checkLogin,name='checkLogin'),
    path('',home),
    path('logout', logout),
    path('deleteFollowing', deleteFollowing,name='deleteFollowing'),
    path('unfollowTopic', unfollowTopic,name='unfollowTopic'),
    path('removeFollower', removeFollower, name='removeFollower'),
    path('post/<int:pid>', viewPost, name='viewPost'),
    path('signup', signup,name='signup'),
    path('viewProfile/<int:uid>',viewProfile,name='viewProfile'),
    path('doUpvote',doUpvote,name='doUpvote'),
    path('doDownvote',doDownvote,name='doDownvote'),
    path('followAjax',followajax,name='followAjax'),
    path('followTopicAjax',followTopicAjax,name='followTopicAjax'),
    path('unfollowUserBtnAjax',unfollowUserBtnAjax,name='unfollowUserBtnAjax'),
    path('followUserBtnAjax',followUserBtnAjax,name='followUserBtnAjax'),
    path('postComment/<int:pid>',postComment,name='postComment'),
    path('downComment',downComment,name='downComment'),
    path('upComment',upComment,name='upComment'),
    path('makePost',makePost,name='makePost'),
    path('fetchNotification',fetchNotification,name='fetchNotification'),
    path('mark_all_as_read',mark_all_as_read,name='mark_all_as_read'),
    path('deletePost',deletePost,name='deletePost'),
    path('saveEditProfile',saveEditProfile,name='saveEditProfile'),
    path('search',search,name='search')

]

urlpatterns += static(settings.MEDIA_URL,document_root=settings.MEDIA_ROOT)


import torch
import torch.nn as nn
import torch.nn.functional as Functional

def square_distance(src, dst):
    B, N, _ = src.shape
    _, M, _ = dst.shape
    dist = -2 * torch.matmul(src.float(), dst.float().permute(0, 2, 1))
    dist += torch.sum(src ** 2, -1).view(B, N, 1)
    dist += torch.sum(dst ** 2, -1).view(B, 1, M)
    return dist

def index_points(points, idx):
    """
    Input:
        points: input points data, [B, N, C]
        idx: sample index data, [B, S]
    Return:
        new_points:, indexed points data, [B, S, C]
    """
    device = points.device
    B = points.shape[0]
    view_shape = list(idx.shape)
    view_shape[1:] = [1] * (len(view_shape) - 1)
    repeat_shape = list(idx.shape)
    repeat_shape[0] = 1
    batch_indices = torch.arange(B, dtype=torch.long).to(device).view(view_shape).repeat(repeat_shape)
    new_points = points[batch_indices, idx, :]
    return new_points

def farthest_point_sample(xyz, npoint):
    """
    Input:
        xyz: pointcloud data, [B, N, 3]
        npoint: number of samples
    Return:
        centroids: sampled pointcloud index, [B, npoint]
    """
    device = xyz.device
    B, N, C = xyz.shape
    centroids = torch.zeros(B, npoint, dtype=torch.long).to(device)
    distance = torch.ones(B, N).to(device) * 1e10
    farthest = torch.randint(0, N, (B,), dtype=torch.long).to(device)
    batch_indices = torch.arange(B, dtype=torch.long).to(device)
    for i in range(npoint):
        centroids[:, i] = farthest
        centroid = xyz[batch_indices, farthest, :].view(B, 1, 3)
        dist = torch.sum((xyz - centroid) ** 2, -1)
        mask = dist < distance
        distance[mask] = dist[mask]
        farthest = torch.max(distance, -1)[1]
    return centroids

def knn(nsample, xyz, new_xyz):
    dist = square_distance(xyz, new_xyz)
    id = torch.topk(dist,k=nsample,dim=1,largest=False)[1]
    id = torch.transpose(id, 1, 2)
    return id

class KNN_dist(nn.Module):
    def __init__(self,k):
        super(KNN_dist, self).__init__()
        self.R = nn.Sequential(
            nn.Linear(10,10),
            nn.LeakyReLU(0.2,inplace=True),
            nn.Linear(10,10),
            nn.LeakyReLU(0.2,inplace=True),
            nn.Linear(10,1),
        )
        self.k=k
    def forward(self,F,vertices):
        id = knn(self.k, vertices, vertices)
        F = index_points(F,id)
        v = index_points(vertices,id)
        v_0 = v[:,:,0,:].unsqueeze(-2).repeat(1,1,self.k,1)
        v_F = torch.cat((v_0, v, v_0-v,torch.norm(v_0-v,dim=-1,p=2).unsqueeze(-1)),-1)
        v_F = self.R(v_F.type(torch.FloatTensor).to('cuda:0'))
        F = torch.mul(v_F, F)
        F = torch.sum(F,-2)
        return F

class View_selector(nn.Module):
    def __init__(self, n_views, sampled_view):
        super(View_selector, self).__init__()
        self.n_views = n_views
        self.s_views = sampled_view
        self.cls = nn.Sequential(
            nn.Linear(512*self.s_views, 256*self.s_views),
            nn.LeakyReLU(0.2),
            nn.Linear(256*self.s_views, 40*self.s_views))
    def forward(self,F,vertices,k):
        id = farthest_point_sample(vertices,self.s_views)
        vertices1 = index_points(vertices,id)
        id_knn = knn(k,vertices,vertices1)
        F = index_points(F,id_knn)
        vertices = index_points(vertices,id_knn)
        F1 = F.transpose(1,2).reshape(F.shape[0],k,self.s_views*F.shape[-1])
        F_score = self.cls(F1).reshape(F.shape[0],k,self.s_views,40).transpose(1,2)
        F1_ = Functional.softmax(F_score,-3)
        F1_ = torch.max(F1_,-1)[0]
        F1_id = torch.argmax(F1_,-1)
        F1_id = Functional.one_hot(F1_id,4).float()
        F1_id_v = F1_id.unsqueeze(-1).repeat(1,1,1,3)
        F1_id_F = F1_id.unsqueeze(-1).repeat(1, 1, 1, 512)
        F_new = torch.mul(F1_id_F,F).sum(-2)
        vertices_new = torch.mul(F1_id_v,vertices).sum(-2)
        return F_new,F_score,vertices_new

class LocalGCN(nn.Module):
    def __init__(self,k,n_views):
        super(LocalGCN,self).__init__()
        self.conv = nn.Sequential(
            nn.Linear(512, 512),
            nn.BatchNorm1d(512),
            nn.LeakyReLU(0.2, inplace=True),
        )
        self.k = k
        self.n_views = n_views
        self.KNN = KNN_dist(k=self.k)
    def forward(self,F,V):
        F = self.KNN(F, V)
        F = F.view(-1, 800)
        F = self.conv(F)
        F = F.view(-1, self.n_views, 756)
        return F

class NonLocalMP(nn.Module):
    def __init__(self,n_view):
        super(NonLocalMP,self).__init__()
        self.n_view=n_view
        self.Relation = nn.Sequential(
            nn.Linear(2 * 512, 512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(512, 512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(512, 512),
            nn.LeakyReLU(0.2, inplace=True),
        )
        self.Fusion = nn.Sequential(
            nn.Linear(2 * 512, 512),
            nn.BatchNorm1d(512),
            nn.LeakyReLU(0.2, inplace=True),
        )
    def forward(self, F):
        F_i = torch.unsqueeze(F, 2)
        F_j = torch.unsqueeze(F, 1)
        F_i = F_i.repeat(1, 1, self.n_view, 1)
        F_j = F_j.repeat(1, self.n_view, 1, 1)
        M = torch.cat((F_i, F_j), 3)
        M = self.Relation(M)
        M = torch.sum(M,-2)
        F = torch.cat((F, M), 2)
        F = F.view(-1, 512 * 2)
        F = self.Fusion(F)
        F = F.view(-1, self.n_view, 512)
        return F


"""Binary tree/non-tree classifier(s)."""
import glob
from os import path

import dask
import maxflow as mf
import numpy as np
import rasterio as rio
from dask import diagnostics
from sklearn import ensemble

from . import pixel_features, pixel_response, settings, utils

__all__ = ["ClassifierTrainer", "Classifier"]

MOORE_NEIGHBORHOOD_ARR = np.array([[0, 0, 0], [0, 0, 1], [1, 1, 1]])

class ClassifierTrainer(object):
    """Train binary tree/non-tree classifier(s) of the pixel features."""

    def __init__(
        self,
        *,
        num_estimators=None,
        sigmas=None,
        num_orientations=None,
        neighborhood=None,
        min_neighborhood_range=None,
        num_neighborhoods=None,
        tree_val=None,
        nontree_val=None,
        **adaboost_kws,
    ):
        """
        Initialize the classifier.

        See the `background <https://bit.ly/2KlCICO>`_ example notebook for details.

        Parameters
        ----------
        num_estimators : int, optional
            The maximum number of estimators at which boosting is terminated. Directly
            passed to the `n_estimators` keyword argument of
            `sklearn.ensemble.AdaBoostClassifier`. If no value is provided, the default
            value set in `settings.CLF_DEFAULT_NUM_ESTIMATORS` will be taken.
        sigmas : list-like, optional
            The list of scale parameters (sigmas) to build the Gaussian filter bank that
            will be used to compute the pixel-level features. The provided argument will
            be passed to the initialization method of the `PixelFeaturesBuilder` class.
            If no value is provided, the default value set in
            `settings.GAUSS_DEFAULT_SIGMAS` will be taken.
        num_orientations : int, optional
            The number of equally-distributed orientations to build the Gaussian filter
            bank that will be used to compute the pixel-level features. The provided
            argument will be passed to the initialization method of the
            `PixelFeaturesBuilder` class. If no value is provided, the default value set
            in `settings.GAUSS_DEFAULT_NUM_ORIENTATIONS` will be taken.
        neighborhood : array-like, optional
            The base neighborhood structure that will be used to compute the entropy
            features. Theprovided argument will be passed to the initialization method
            of the `PixelFeaturesBuilder` class. If no value is provided, a square with
            a side size of `2 * min_neighborhood_range + 1` will be used.
        min_neighborhood_range : int, optional
            The range (i.e., the square radius) of the smallest neigbhorhood window that
            will be used to compute the entropy features. The provided argument will be
            passed to the initialization method of the `PixelFeaturesBuilder` class. If
            no value is provided, the default value set in
            `settings.ENTROPY_DEFAULT_MIN_NEIGHBORHOOD_RANGE` will be taken.
        num_neighborhoods : int, optional
            The number of neigbhorhood windows (whose size follows a geometric
            progression starting at `min_neighborhood_range`) that will be used to
            compute the entropy features. The provided argument will be passed to the
            initialization method of the `PixelFeaturesBuilder` class. If no value is
            provided, the default value set in
            `settings.ENTROPY_DEFAULT_NUM_NEIGHBORHOODS` will be taken.
        tree_val : int, optional
            The value that designates tree pixels in the response images. The provided
            argument will be passed to the initialization method of the
            `PixelResponseBuilder` class. If no value is provided, the default value set
            in `settings.RESPONSE_DEFAULT_TREE_VAL` will be taken.
        nontree_val : int, optional
            The value that designates non-tree pixels in the response images. The
            provided argument will be passed to the initialization method of the
            `PixelResponseBuilder` class. If no value is provided, the default value set
            in `settings.RESPONSE_DEFAULT_NONTREE_VAL` will be taken.
        adaboost_kws : key-value pairings, optional
            Keyword arguments that will be passed to
            `sklearn.ensemble.AdaBoostClassifier`.
        """
        super(ClassifierTrainer, self).__init__()

        if num_estimators is None:
            num_estimators = settings.CLF_DEFAULT_NUM_ESTIMATORS
        self.num_estimators = num_estimators

        self.pixel_features_builder_kws = dict(
            sigmas=sigmas,
            num_orientations=num_orientations,
            neighborhood=neighborhood,
            min_neighborhood_range=min_neighborhood_range,
            num_neighborhoods=num_neighborhoods,
        )
        self.pixel_response_builder_kws = dict(
            tree_val=tree_val, nontree_val=nontree_val
        )
        self.adaboost_kws = adaboost_kws

    def train_classifier(
        self,
        *,
        split_df=None,
        response_img_dir=None,
        img_filepaths=None,
        response_img_filepaths=None,
        img_dir=None,
        img_filename_pattern=None,
        method=None,
        img_cluster=None,
    ):
        """
        Train a classifier.

        See the `background <https://bit.ly/2KlCICO>`_ example notebook for more
        details.

        Parameters
        ----------
        split_df : pandas DataFrame, optional
            Data frame with the train/test split.
        response_img_dir : str representing path to a directory, optional
            Path to the directory where the response tiles are located. Required if
            providing `split_df`. Otherwise `response_img_dir` might either be ignored
            if providing `response_img_filepaths`, or be used as the directory where the
            images whose filename matches `img_filename_pattern` are to be located.
        img_filepaths : list-like, optional
            List of paths to the input tiles whose features will be used to train the
            classifier. Ignored if `split_df` is provided.
        response_img_filepaths : list-like, optional
            List of paths to the binary response tiles that will be used to train the
            classifier. Ignored if `split_df` is provided.
        img_dir : str representing path to a directory, optional
            Path to the directory where the images whose filename matches
            `img_filename_pattern` are to be located. Ignored if `split_df` or
            `img_filepaths` is provided.
        img_filename_pattern : str representing a file-name pattern, optional
            Filename pattern to be matched in order to obtain the list of images. If no
            value is provided, the default value set in
            `settings.IMG_DEFAULT_FILENAME_PATTERN` will be taken. Ignored if `split_df`
            or `img_filepaths` is provided.
        method : {'cluster-I', 'cluster-II'}, optional
            Method used in the train/test split.
        img_cluster : int, optional
            The label of the cluster of tiles. Only used if `method` is 'cluster-II'.

        Returns
        -------
        clf : scikit-learn AdaBoostClassifier
            The trained classifier
        """
        if split_df is None and response_img_filepaths is None:
            # this is the only case that needs argument tweaking: otherwise, if we pass
            # `img_filepaths`/`img_dir` to `build_features` and `response_img_dir` to
            # `build_response`, the latter would build a response with all the image
            # files in `response_img_dir`. Instead, we need to build the response only
            # for the files speficied in `img_filepaths`/`img_dir`
            if img_filepaths is None:
                # TODO: this is copied from `build_features` - ideally, we should DRY it
                if img_filename_pattern is None:
                    img_filename_pattern = settings.IMG_DEFAULT_FILENAME_PATTERN
                if img_dir is None:
                    raise ValueError(
                        "Either `split_df`, `img_filepaths` or `img_dir` must "
                        "be provided"
                    )
                img_filepaths = glob.glob(path.join(img_dir, img_filename_pattern))

            response_img_filepaths = [
                path.join(response_img_dir, path.basename(img_filepath))
                for img_filepath in img_filepaths
            ]

        X = pixel_features.PixelFeaturesBuilder(
            **self.pixel_features_builder_kws
        ).build_features(
            split_df=split_df,
            img_filepaths=img_filepaths,
            img_dir=img_dir,
            img_filename_pattern=img_filename_pattern,
            method=method,
            img_cluster=img_cluster,
        )

        y = pixel_response.PixelResponseBuilder(
            **self.pixel_response_builder_kws
        ).build_response(
            split_df=split_df,
            response_img_dir=response_img_dir,
            response_img_filepaths=response_img_filepaths,
            img_filename_pattern=img_filename_pattern,
            method=method,
            img_cluster=img_cluster,
        )

        clf = ensemble.AdaBoostClassifier(
            n_estimators=self.num_estimators, **self.adaboost_kws
        )
        clf.fit(X, y)

        return clf

    def train_classifiers(self, split_df, response_img_dir):
        """
        Train a classifier for each first-level cluster in `split_df`.

        See the `background <https://bit.ly/2KlCICO>`_ example notebook for more
        details.

        Parameters
        ----------
        split_df : pandas DataFrame
            Data frame with the train/test split, which must have an `img_cluster`.
            column with the first-level cluster labels.
        response_img_dir : str representing path to a directory
            Path to the directory where the response tiles are located.

        Returns
        -------
        clf_dict : dictionary
            Dictionary mapping a scikit-learn AdaBoostClassifier to each first-level
            cluster label.
        """
        if "img_cluster" not in split_df:
            raise ValueError(
                "`split_df` must have an 'img_cluster' column ('cluster-II'). "
                "For 'cluster-I', use `train_classifier`."
            )

        clfs_lazy = {}
        for img_cluster, _ in split_df.groupby("img_cluster"):
            clfs_lazy[img_cluster] = dask.delayed(self.train_classifier)(
                split_df=split_df,
                response_img_dir=response_img_dir,
                method="cluster-II",
                img_cluster=img_cluster,
            )

        with diagnostics.ProgressBar():
            clfs_dict = dask.compute(clfs_lazy)[0]

        return clfs_dict

class Classifier(object):
    """Use trained classifier(s) to predict tree pixels."""

    def __init__(
        self,
        *,
        tree_val=None,
        nontree_val=None,
        refine=None,
        refine_beta=None,
        refine_int_rescale=None,
        **pixel_features_builder_kws,
    ):
        """
        Initialize the classifier instance.

        See the `background <https://bit.ly/2KlCICO>`_ example notebook for more
        details.

        Parameters
        ----------
        tree_val : int, optional
            Label used to denote tree pixels in the predicted images. If no value is
            provided, the default value set in `settings.CLF_DEFAULT_TREE_VAL` will be
            taken.
        nontree_val : int, optional
            Label used to denote non-tree pixels in the predicted images. If no value is
            provided, the default value set in `settings.CLF_DEFAULT_NONTREE_VAL` will
            be taken.
        refine : bool, optional
            Whether the pixel-level classification should be refined by optimizing the
            consistence between neighboring pixels. If no value is provided, the default
            value set in `settings.CLF_DEFAULT_REFINE` will be taken.
        refine_beta : int, optional
            Parameter of the refinement procedure that controls the smoothness of the
            labelling. Larger values lead to smoother shapes.  If no value is provided,
            the default value set in `settings.CLF_DEFAULT_REFINE_BETA` will be taken.
        refine_int_rescale : int, optional
            Parameter of the refinement procedure that controls the precision of the
            transformation of float to integer edge weights, required for the employed
            graph cuts algorithm. Larger values lead to greater precision. If no value
            is provided, the default value set in
            `settings.CLF_DEFAULT_REFINE_INT_RESCALE` will be taken.
        pixel_features_builder_kws : dict, optional
            Keyword arguments that will be passed to `detectree.PixelFeaturesBuilder`,
            which customize how the pixel features are built.
        """
        super(Classifier, self).__init__()

        if tree_val is None:
            tree_val = settings.CLF_DEFAULT_TREE_VAL
        if nontree_val is None:
            nontree_val = settings.CLF_DEFAULT_NONTREE_VAL
        if refine is None:
            refine = settings.CLF_DEFAULT_REFINE
        if refine_beta is None:
            refine_beta = settings.CLF_DEFAULT_REFINE_BETA
        if refine_int_rescale is None:
            refine_int_rescale = settings.CLF_DEFAULT_REFINE_INT_RESCALE

        self.tree_val = tree_val
        self.nontree_val = nontree_val
        self.refine = refine
        self.refine_beta = refine_beta
        self.refine_int_rescale = refine_int_rescale

        self.pixel_features_builder_kws = pixel_features_builder_kws

    def classify_img(self, img_filepath, clf, output_filepath=None):
        """
        Use a trained classifier to predict tree pixels in an image.

        Optionally dump the predicted tree/non-tree image to `output_filepath`.

        Parameters
        ----------
        img_filepath : str, file object or pathlib.Path object
            Path to a file, URI, file object opened in binary ('rb') mode, or a Path
            object representing the image to be classified. The value will be passed to
            `rasterio.open`.
        clf : scikit-learn AdaBoostClassifier
            Trained classifier.
        output_filepath : str, file object or pathlib.Path object, optional
            Path to a file, URI, file object opened in binary ('rb') mode, or a Path
            object representing where the predicted image is to be dumped. The value
            will be passed to `rasterio.open` in 'write' mode.

        Returns
        -------
        y_pred : numpy ndarray
            Array with the pixel responses.
        """
        # ACHTUNG: Note that we do not use keyword-only arguments in this method because
        # `output_filepath` works as the only "optional" argument
        src = rio.open(img_filepath)
        img_shape = src.shape

        X = pixel_features.PixelFeaturesBuilder(
            **self.pixel_features_builder_kws
        ).build_features_from_filepath(img_filepath)

        if not self.refine:
            y_pred = clf.predict(X).reshape(img_shape)
        else:
            p_nontree, p_tree = np.hsplit(clf.predict_proba(X), 2)
            g = mf.Graph[int]()
            node_ids = g.add_grid_nodes(img_shape)
            P_nontree = p_nontree.reshape(img_shape)
            P_tree = p_tree.reshape(img_shape)

            # The AdaBoost probabilities are floats between 0 and 1, and the graph cuts
            # algorithm requires an integer representation.  Therefore, we will multiply
            # the probabilities by an arbitrary large number and then transform the
            # result to integers. For instance, we could use a `refine_int_rescale` of
            # `100` so that the probabilities are rescaled into integers between 0 and
            # 100 like percentages). The larger `refine_int_rescale`, the greater the
            # precision.
            # ACHTUNG: the data term when the pixel is a tree is `log(1 - P_tree)`,
            # i.e., `log(P_nontree)`, so the two lines below are correct
            D_tree = (self.refine_int_rescale * np.log(P_nontree)).astype(int)
            D_nontree = (self.refine_int_rescale * np.log(P_tree)).astype(int)
            # TODO: option to choose Moore/Von Neumann neighborhood?
            g.add_grid_edges(
                node_ids, self.refine_beta, structure=MOORE_NEIGHBORHOOD_ARR
            )
            g.add_grid_tedges(node_ids, D_tree, D_nontree)
            g.maxflow()
            # y_pred = g.get_grid_segments(node_ids)
            # transform boolean `g.get_grid_segments(node_ids)` to an array of
            # `self.tree_val` and `self.nontree_val`
            y_pred = np.full(img_shape, self.nontree_val)
            y_pred[g.get_grid_segments(node_ids)] = self.tree_val

        # TODO: make the profile of output rasters more customizable (e.g., via the
        # `settings` module)
        # output_filepath = path.join(output_dir,
        #                             f"tile_{tile_start}-{tile_end}.tif")
        if output_filepath is not None:
            with rio.open(
                output_filepath,
                "w",
                driver="GTiff",
                width=y_pred.shape[1],
                height=y_pred.shape[0],
                count=1,
                dtype=np.uint8,
                nodata=self.nontree_val,
                crs=src.crs,
                transform=src.transform,
            ) as dst:
                dst.write(y_pred.astype(np.uint8), 1)

        src.close()
        return y_pred

    def _classify_imgs(self, img_filepaths, clf, output_dir):
        pred_imgs_lazy = []
        pred_img_filepaths = []
        for img_filepath in img_filepaths:
            # filename, ext = path.splitext(path.basename(img_filepath))
            # pred_img_filepath = path.join(
            #     output_dir, f"{filename}-pred{ext}")
            pred_img_filepath = path.join(output_dir, path.basename(img_filepath))
            pred_imgs_lazy.append(
                dask.delayed(self.classify_img)(img_filepath, clf, pred_img_filepath)
            )
            pred_img_filepaths.append(pred_img_filepath)

        with diagnostics.ProgressBar():
            dask.compute(*pred_imgs_lazy)

        return pred_img_filepaths

    def classify_imgs(
        self,
        split_df,
        output_dir,
        *,
        clf=None,
        clf_dict=None,
        method=None,
        img_cluster=None,
    ):
        """
        Use trained classifier(s) to predict tree pixels in multiple images.

        Use `clf` or `clf_dict` for the classifier(s) depending on the train/ test split
        method, and dump the predicted tree/non-tree images to `output_dir`. See the
        `background <https://bit.ly/2KlCICO>`_ example notebook for more details.

        Parameters
        ----------
        split_df : pandas DataFrame, optional
            Data frame with the train/test split.
        output_dir : str or pathlib.Path object
            Path to the directory where the predicted images are to be dumped.
        clf : scikit-learn AdaBoostClassifier
            Trained classifier.
        clf_dict : dictionary
            Dictionary mapping a trained scikit-learn AdaBoostClassifier to each
            first-level cluster label.
        method : {'cluster-I', 'cluster-II'}, optional
            Method used in the train/test split.
        img_cluster : int, optional
            The label of the cluster of tiles. Only used if `method` is 'cluster-II'.

        Returns
        -------
        pred_imgs : list or dict
            File paths of the dumped tiles.
        """
        if method is None:
            if "img_cluster" in split_df:
                method = "cluster-II"
            else:
                method = "cluster-I"

        if method == "cluster-I":
            if clf is None:
                raise ValueError("If using 'cluster-I' method, `clf` must be provided")
            return self._classify_imgs(
                split_df[~split_df["train"]]["img_filepath"], clf, output_dir
            )
        else:
            if img_cluster is not None:
                if clf is None:
                    if clf_dict is not None:
                        clf = clf_dict[img_cluster]
                    else:
                        raise ValueError("Either `clf` or `clf_dict` must be provided")

                return self._classify_imgs(
                    utils.get_img_filepaths(split_df, img_cluster, False),
                    clf,
                    output_dir,
                )

            if clf_dict is None:
                raise ValueError(
                    "If using 'cluster-II' method and not providing `img_cluster`, "
                    "`clf_dict` must be provided"
                )
            pred_imgs = {}
            for img_cluster, _ in split_df.groupby("img_cluster"):
                pred_imgs[img_cluster] = self._classify_imgs(
                    utils.get_img_filepaths(split_df, img_cluster, False),
                    clf_dict[img_cluster],
                    output_dir,
                )

            return pred_imgs


import cocotb
from cocotb.clock import Clock
from cocotb.triggers import RisingEdge, FallingEdge, Timer, ClockCycles

@cocotb.test()
async def test_hex_sr(dut):
    dut._log.info("start")
    clock = Clock(dut.clk, 10, units="us")
    cocotb.start_soon(clock.start())

    cycles = 512
    sr_len = 100

    dut.recirc.value = 0
    for cycle in range(cycles):
        dut.sr_in = cycle % 64
        await RisingEdge(dut.clk)
        if cycle >= sr_len:
            v = int(dut.sr_out)
            assert v == (cycle - sr_len) % 64
            # output matches input with 40 cycle latency
        await FallingEdge(dut.clk)

    dut.recirc.value = 1
    for cycle in range(cycles, 2 * cycles):
        dut.sr_in = cycle % 64
        await RisingEdge(dut.clk)
        v = int(dut.sr_out)
        if cycle >= cycles + sr_len:
            v = int(dut.sr_out)
            assert v == (cycle - sr_len) # 64
        await FallingEdge(dut.clk)


"""
[베스트 앨범]
https://school.programmers.co.kr/learn/courses/30/lessons/42579

스트리밍 사이트에서 장르 별로 가장 많이 재생된 노래를 두 개씩 모아 베스트 앨범을 출시
노래를 수록하는 기준
- 속한 노래가 많이 재생된 장르를 먼저 수록합니다.
- 장르 내에서 많이 재생된 노래를 먼저 수록합니다.
- 장르 내에서 재생 횟수가 같은 노래 중에서는 고유 번호가 낮은 노래를 먼저 수록합니다.

제한 사항
- genres와 plays의 길이는 동일
- 장르 종류 < 100개
- 장르에 속한 곡이 하나라면, 하나의 곡만 선택합니다.
- 모든 장르는 재생된 횟수가 다르다.
"""
from typing import List
from collections import defaultdict

def solution(genres:List, plays:List) -> List:
    answer = []

    # key에 대한 기본 데이터 형태가 정해지면 defaultdict를 사용하기
    total_info = defaultdict(list)
    songs_per_genre = defaultdict(int)

    for i, (genre, play_count) in enumerate(zip(genres, plays)):
        # {'classic': [(0, 500), (2, 150), (3, 800)], 'pop': [(1, 600), (4, 2500)]}
        total_info[genre].append((i, play_count))

        # {'classic': 1450, 'pop': 3100}
        songs_per_genre[genre] += play_count

    # item의 1 인덱스에 있는 게 재생 횟수이므로, 이를 기준으로 내림차순으로 정렬한다.
    # 정렬하여 재생횟수가 제일 많은 것부터 플레이 리스트에 추가한다.
    for genre, _ in sorted(songs_per_genre.items(), key=lambda x:x[1], reverse=True):
        for i, _ in sorted(total_info[genre], key=lambda x:x[1], reverse=True)[:2]:
            answer.append(i)

    return answer

if __name__ == "__main__":
    genres = ["classic", "pop", "classic", "classic", "pop"]
    plays = [500, 600, 150, 800, 2500]
    print(solution(genres, plays))


import torch
import torch.nn as nn
import functools
from loss import *

class Res_Enhancement(nn.Module):
    def __init__(self):
        super(Res_Enhancement, self).__init__()

        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1),
            nn.ReLU(inplace=True))
        self.conv2_0 = nn.Sequential(
            nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, padding=1),
            nn.BatchNorm2d(16),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, padding=1),
            nn.BatchNorm2d(16))

        self.conv2_1 = nn.Sequential(
            nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, padding=1),
            nn.BatchNorm2d(16),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, padding=1),
            nn.BatchNorm2d(16))

        self.conv2_2 = nn.Sequential(
            nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, padding=1),
            nn.BatchNorm2d(16),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, padding=1),
            nn.BatchNorm2d(16))

        self.delta = nn.Sequential(
            nn.Conv2d(in_channels=16, out_channels=3, kernel_size=1, padding=0),
            nn.Sigmoid())

        self.en_loss = LossFunction()

    def forward(self, x):
        block1 = self.conv1(x)
        block2 = self.conv2_0(block1)
        block2 = block1 + block2
        block3 = self.conv2_1(block2)
        block3 = block2 + block3
        block4 = self.conv2_2(block3)
        block4 = block3 + block4
        delta = self.delta(block4)
        illu = delta + x

        illu = torch.clamp(illu, 0.001, 1.0)
        enhance = x / illu
        enhance = torch.clamp(enhance, 0.0, 1.0)
        return enhance, illu

    def loss(self, x):
        illu, enhance = self(x)
        loss = self.en_loss(x, illu)
        return loss