from transformers import AutoTokenizer, AutoModelForMaskedLM,  RobertaConfig , RobertaTokenizer,RobertaForMaskedLM, DataCollatorForLanguageModeling, LineByLineTextDataset, Trainer, TrainingArguments

from pathlib import Path
from tokenizers import ByteLevelBPETokenizer
from tokenizers.implementations import ByteLevelBPETokenizer
from tokenizers.processors import BertProcessing
import torch
from torchinfo import summary

import os

paths = [str(x) for x in Path(".").glob("**/el_*.txt")]
print(paths)
# Initialize a tokenizer
tokenizer = ByteLevelBPETokenizer()
# Customize training
tokenizer.train(files=paths, vocab_size=52_000, min_frequency=2,
special_tokens=[
"<s>",
"<pad>",
"</s>",
"<unk>",
"<mask>",
])

dir_path = os.getcwd()
token_dir = os.path.join(dir_path, 'QuijoBERT')

if not os.path.exists(token_dir):
  os.makedirs(token_dir)
tokenizer.save_model('QuijoBERT')

tokenizer = ByteLevelBPETokenizer(
"./QuijoBERT/vocab.json",
"./QuijoBERT/merges.txt",
)

tokenizer._tokenizer.post_processor = BertProcessing(
("</s>", tokenizer.token_to_id("</s>")),
("<s>", tokenizer.token_to_id("<s>")),
)
tokenizer.enable_truncation(max_length=512)

config = RobertaConfig(
  vocab_size=52_000,
  max_position_embeddings=514,
  num_attention_heads=12,
  num_hidden_layers=6,
  type_vocab_size=1,
  )

"""# Step 8: Re-creating the Tokenizer in Transformers"""

tokenizer = RobertaTokenizer.from_pretrained("./QuijoBERT", max_length=512)

#Initializing a Model

model = RobertaForMaskedLM(config=config)
#In case we want to recover the after a crash
#model = RobertaForMaskedLM.from_pretrained("./QuijoBERT/Checkpoint-xxxxx")

#Tensorflow
print(model)
#Pytorch
summary(model)

dataset = LineByLineTextDataset(
  tokenizer=tokenizer,
  file_path="./el_quijote.txt",
  block_size=128,
  )

#Defining a Data Collator

data_collator = DataCollatorForLanguageModeling(
  tokenizer=tokenizer, mlm=True, mlm_probability=0.15
)

# Initializing the Trainer Object
training_args = TrainingArguments(
  output_dir="./QuijoBERT",
  overwrite_output_dir=True,
  num_train_epochs=1,
  per_device_train_batch_size=64,
  save_steps=1000,
  save_total_limit=2,
  )
trainer = Trainer(
  model=model,
  args=training_args,
  data_collator=data_collator,
  train_dataset=dataset,
)

#Training the Model
print('aqui')
trainer.train()
trainer.save_model("./QuijoBERT")

#Saving the Final Model(+tokenizer + config) to disk
trainer.save_model("./QuijoBERT")


import torch
import torchvision.models

shape = (1,3,64,64)
model = torchvision.models.resnet18(pretrained=True)
data = torch.rand(shape)
labels = torch.rand(1,1000)

# forward

prediction = model(data)

loss = (prediction - labels).sum()

loss.backward()

optim = torch.optim.SGD(model.parameters(),lr=1e-2,momentum=0.9)

optim.step()


#!/home/yin/opt/bin/python3

import numpy as np
from myvasp import vasp_func as vf

def main():
    jobn, Etot, Eent, pres = vf.vasp_read_post_data()

    if Etot.shape[0] < 5.9:
        import sys
        sys.exit("\n==> ABORT: insufficient data \n" )

    pres0 = vf.read_pressure('../y_full_relax/OUTCAR')
    temp = np.vstack( (pres, pres0) )

    s=temp.copy()
    s[:,3] = temp[:,4]
    s[:,4] = temp[:,5]
    s[:,5] = temp[:,3]

    s = s.T *(-0.1)   # stress data, 6*7, [GPa]

    s_d, Cij_d = calc_direct_Cij(s)
    fitres, C14, Cij_cubic, Cij_hcp = calc_fitting_Cij(s)
    write_output(s, s_d, Cij_d, fitres, C14, Cij_cubic, Cij_hcp)

def calc_direct_Cij(s):
    s_d = ( s[:,0:6].copy().T - s[:,-1].copy() ).T
    Cij_d = s_d /0.002
    return s_d, Cij_d

def calc_fitting_Cij(s):
    from sympy import MatrixSymbol, Matrix, eye, zeros, diff, solve

    Cij = zeros(6)
    Stress0 = zeros(6, 7)

    c = MatrixSymbol('c', 6, 6)
    s0 = MatrixSymbol('s0', 6, 1)

    for i in np.arange(0, 6, 1):
        for j in np.arange(i, 6, 1):
            Cij[i,j] = c[i,j]
            Cij[j,i] = c[i,j]

        for k in np.arange(7):
            Stress0[i,k] = s0[i,0]
    # print( Cij, Stress0 )

    Strain = eye(6)*0.002
    Strain = Matrix([Strain, zeros(1, 6)])
    Strain = Strain.T
    # print(Strain)

    Stress = Cij * Strain + Stress0

    err=Stress - s
    # print(err)

    y2=0
    for i in np.arange(err.shape[0]):
        for j in np.arange(err.shape[1]):
            y2 = y2 + (err[i,j])**2

    # print('==> y2:')
    # print(y2)

    myeqn = zeros(27, 1)
    myvar = zeros(27, 1)
    k=-1
    for i in np.arange(6):
        for j in np.arange(i, 6, 1):
            k=k+1
            myeqn[k,0] = diff(y2, c[i,j])
            myvar[k,0] = c[i,j]

    for i in np.arange(6):
        k=k+1
        myeqn[k,0] = diff(y2, s0[i,0])
        myvar[k,0] = s0[i,0]

    fitres = solve(myeqn, myvar)

    C14=[]
    for i in np.arange(3):
        for j in np.arange(3, 6, 1):
            C14.append( fitres[c[i, j]] )
    C14 = np.array(C14, dtype=np.float64 )

    # for cubic
    Cij_cubic = [
    [ fitres[c[0, 0]] , fitres[c[1, 1]] , fitres[c[2, 2]] ],
    [ fitres[c[0, 1]] , fitres[c[0, 2]] , fitres[c[1, 2]] ],
    [ fitres[c[3, 3]] , fitres[c[4, 4]] , fitres[c[5, 5]] ],
    ]

    # for hcp
    Cij_hcp = [
    [ fitres[c[0, 0]] , fitres[c[1, 1]]  ],
    [ fitres[c[0, 1]]  ],
    [ fitres[c[0, 2]] , fitres[c[1, 2]]  ],
    [ fitres[c[2, 2]]  ],
    [ fitres[c[3, 3]] , fitres[c[4, 4]]  ],
    [ fitres[c[5, 5]]  ],
    ]

    return fitres, C14, Cij_cubic, Cij_hcp

def write_output(s, s_d, Cij_d, fitres, C14, Cij_cubic, Cij_hcp):

    from sympy import MatrixSymbol
    c = MatrixSymbol('c', 6, 6)
    s0 = MatrixSymbol('s0', 6, 1)

    f = open("y_post_cij_stress.txt","w+")

    f.write("# Cij from stress-strain method: \n" )

    f.write("\n# stress data (GPa): \n"  )
    for i in np.arange(s.shape[0]):
        for j in np.arange(s.shape[1]):
            f.write("%10.6f " %(s[i,j]) )
        f.write(" \n")

    f.write("\n# stress_d data (GPa): \n"  )
    for i in np.arange(s_d.shape[0]):
        for j in np.arange(s_d.shape[1]):
            f.write("%10.6f " %(s_d[i,j]) )
        f.write(" \n")

    f.write("\n# Cij_d, direct results: \n"  )
    for i in np.arange(Cij_d.shape[0]):
        for j in np.arange(Cij_d.shape[1]):
            f.write("%10.2f " %(Cij_d[i,j]) )
        f.write(" \n")

    f.write("\n# Cij, fitting results: \n"  )
    for i in np.arange(6):
        for j in np.arange(6):
            if j<i:
                f.write("%10s " %('*') )
            else:
                f.write("%10.2f " %(fitres[c[i, j]]) )
        f.write(" \n")

    f.write("\n# Stress0, fitting results: \n"  )
    for i in np.arange(6):
        f.write("%10.6f " %(fitres[s0[i, 0]]) )
        f.write(" \n")

    f.write("\n# mean and std: \n"  )
    f.write("# C14: \n"  )
    f.write("%10.2f %10.2f \n" %( C14.mean(), C14.std()) )

    f.write("\n# if cubic, C11, C12, C44: \n"  )
    cij_mean=np.array([])
    for i in np.arange(len(Cij_cubic)):
        temp=np.array(Cij_cubic[i],  dtype=np.float64 )
        cij_mean = np.append(cij_mean, temp.mean() )
        f.write("%10.2f %10.2f \n" %( temp.mean(), temp.std()) )

    f.write("\n# B_fcc: \n"  )
    f.write("%10.2f \n\n" %( (cij_mean[0] + 2*cij_mean[1])/3 ))

    f.write("\n# if hcp, C11, C12, C13, C33, C44, (C66): \n"  )
    cij_mean=np.array([])
    for i in np.arange(len(Cij_hcp)):
        temp=np.array(Cij_hcp[i],  dtype=np.float64 )
        cij_mean = np.append(cij_mean, temp.mean() )
        f.write("%10.2f %10.2f \n" %( temp.mean(), temp.std()) )

    f.write("\n# C66-(C11-C12)/2: \n"  )
    f.write("%10.2f \n" %( cij_mean[5] - ( cij_mean[0] - cij_mean[1] )/2 ))

    f.write("\n# B_hcp: \n"  )
    c11=cij_mean[0]
    c12=cij_mean[1]
    c13=cij_mean[2]
    c33=cij_mean[3]
    c44=cij_mean[4]

    f.write("%10.2f \n\n" %( (c11*c33 + c12*c33 - 2*c13**2)/(c11 + c12 - 4*c13 + 2*c33) ))

    f.write("# Transverse isotropy \n")

    from myalloy import calc_elastic_constant as ce
    E_x, E_z, nu_xy, nu_xz, mu_xz = ce.calc_transverse_isotropy( cij_mean[0:-1] )

    f.write('\n%12s %12s %12s %12s %12s \n'
        %('E_x', 'E_z', 'nu_xy', 'nu_xz', 'mu_xz') )

    f.write('%12.4f %12.4f %12.4f %12.4f %12.4f \n' \
        %( E_x,   E_z,   nu_xy,   nu_xz,   mu_xz ) )

    f.write('\n%25s %25s \n'
        %('E_x/(1-nu_xy)', 'nu_xz/(1-nu_xy)' ) )

    f.write('%25.4f %25.4f \n' \
        %( E_x/(1-nu_xy),   nu_xz/(1-nu_xy)  ) )

    f.write('\n%25s \n'
        %('E_x/(1-nu_xy**2)'  ) )

    f.write('%25.4f \n' \
        %( E_x/(1-nu_xy**2)   ) )

    f.close()

main()


import time,board,busio
import numpy as np
import adafruit_mlx90640
import matplotlib.pyplot as plt
import pandas as pd
# import keyboard
from tempfile import TemporaryFile

# Set up MLX90640
i2c = busio.I2C(board.SCL, board.SDA, frequency=400000) # setup I2C
mlx = adafruit_mlx90640.MLX90640(i2c) # begin MLX90640 with I2C comm
mlx.refresh_rate = adafruit_mlx90640.RefreshRate.REFRESH_2_HZ # set refresh rate
mlx_shape = (24,32)

# setup the figure for plotting
plt.ion() # enables interactive plotting
fig,ax = plt.subplots(figsize=(7,2))
therm1 = ax.imshow(np.zeros(mlx_shape),vmin=0,vmax=60) #start plot with zeros
cbar = fig.colorbar(therm1) # setup colorbar for temps
cbar.set_label('Temperature [$^{\circ}$C]',fontsize=14) # colorbar label

df = pd.DataFrame(columns=['Timestamp'])
numpyfile = TemporaryFile()
pixel_matrices = []

t_array = []
i = 0

while True:
    t1 = time.monotonic()
    try:
        frame = np.zeros((24*32,)) # setup array for storing all 768 temperatures
        data_array = mlx.getFrame(frame) # read MLX temperatures into frame var
        data_array = (np.reshape(frame,mlx_shape)) # reshape to 24x32
        therm1.set_data(np.fliplr(data_array)) # flip left to right

        # Append pixel_matrix and timestamp values
        pixel_matrices.append(data_array)
        df.at[i,0] = t1
        i = i + 1

        # Save thermal image to directory
        therm1.set_clim(vmin=np.min(data_array),vmax=np.max(data_array)) # set bounds
        cbar.update_normal(therm1) # update colorbar range
        plt.pause(0.001) # required

        fig.savefig(r"/home/pi/Thermcam/Pixel_data/data_0/" + str(t1) + r".png", dpi=300,facecolor='#FCFCFC',
                    bbox_inches='tight') # comment out to speed up
        t_array.append(time.monotonic()-t1)

        # print('Sample Rate: {0:2.1f}fps'.format(len(t_array)/np.sum(t_array)))

    except KeyboardInterrupt:
        numpyfile = str(t1)
        pixel_array = np.stack(pixel_matrices)
        np.save(r"/home/pi/Thermcam/Pixel_data/data_0/" + numpyfile, pixel_array, allow_pickle = True)
        df.to_csv(r"/home/pi/Thermcam/Pixel_data/data_0/" + str(t1) + ".csv")
        break


import nltk
from nltk.tokenize import TweetTokenizer
from nltk.corpus import stopwords
from nltk.corpus import opinion_lexicon
import string
from nltk.stem.porter import *
from sklearn.feature_extraction.text import CountVectorizer
import math
import csv

class TweetEnricher:
    """
    Class that adds features to tweets collected
    """
    def __init__(self):
        self.tokenizer = TweetTokenizer()
        self.stemmer = PorterStemmer()
        self.stopset = set(stopwords.words('english'))
        self.negative_opinions = opinion_lexicon.negative()
        self.positive_opinions = opinion_lexicon.positive()
        self.brexit_keywords = [line.rstrip('\n') for line in open('../Data/Lists/BrexitKeywords')]
        self.vulgar_words = [line.rstrip('\n').lower() for line in open('../Data/Lists/VulgarWordsList')]
        self.twitter_jargons = [line.rstrip('\n') for line in open('../Data/Lists/TwitterSlangsAndAbbreviations')]
        self.web_abbreviations = [line.rstrip('\n').lower() for line in open('../Data/Lists/WebAcronymns')]
        self.emoticons_list = [line.rstrip('\n') for line in open('../Data/Lists/EmojiList')]
        self.pos_emoticons_list = [line.rstrip('\n') for line in open('../Data/Lists/PositiveEmojiList')]
        self.neg_emoticons_list = [line.rstrip('\n') for line in open('../Data/Lists/NegativeEmojiList')]
        self.first_person_pronouns = [line.rstrip('\n') for line in open('../Data/Lists/FirstPersonPronouns')]
        self.speech_act_verbs = [line.rstrip('\n') for line in open('../Data/Lists/StemmedSpeechActVerbs')]
        self.trusted_domains = [line.rstrip('\n') for line in open('../Data/Lists/TrustedDomains')]
        self.verb_tags = ['VB', 'VBD', 'VBG', 'VBN', 'VBP', 'VBZ']
        self.n_gram_count_matrix = {}
        self.vectorizer = CountVectorizer(ngram_range=(1, 3), tokenizer=self.tokenizer.tokenize,
                                          stop_words=list(self.stopset)
                                                     + self.web_abbreviations
                                                     + list(string.punctuation)
                                                     + ["…", "...", "..", ")", "(", "-->", "->", ">>", "#", "RT", "@"])
        self.vectorizer_unigram = CountVectorizer(ngram_range=(1, 1), tokenizer=self.tokenizer.tokenize,
                                                  stop_words=list(self.stopset)
                                                             + self.web_abbreviations
                                                             + list(string.punctuation))
        self.positive_ops = [x.lower() for x in self.positive_opinions]
        self.negative_ops = [x.lower() for x in self.negative_opinions]

    def tweetFeatures(self):
        self.tweet_features = ["VulgarWords", "Emoticons", "Interrogation", "Exclamation", "Abbreviations",
                               "TwitterJargons", "#", "# position", "@", "@ position", "RT", "RT position", "Link"]
        for w in self.speech_act_verbs:
            self.tweet_features.append(w)
        self.tweet_features.append(["NegativeOpnions", "PositiveOpnions","NumbersPresent","Non-Ascii Characters","GoodLink"])
        self.basic_features = self.tweet_features
        for w in self.n_gram_count_matrix:
            self.tweet_features.append(w)
        return self.basic_features, self.tweet_features, self.n_gram_count_matrix

    def tokenize(self,tweet):
        """
        Tokens created using nltk tokenizer
        :return: tokens from tweet
        """
        return self.tokenizer.tokenize(tweet)

    def removeStopWords(self,tokens):
        """
        Removes stop words from tokens
        :return: tokens without stopwords
        """
        tokens_without_stopwords = [w for w in tokens if w not in self.stopset]
        return tokens_without_stopwords

    def hasNegativeOpinions(self,tokens):
        """
        Checks if negative opinions from nltk corpus present
        :return: 1 if negative opinions present. 0 otherwise
        """
        count = 0
        negative = 0
        negative_ops = [x.lower() for x in self.negative_opinions]
        for w in tokens:
            if w.lower() in negative_ops:
                count += 1
        if count > 0:
            negative = 1
        return count, negative

    def hasPositiveOpinions(self,tokens):
        '''
        Checks if positive opnions present
        :return: 1 if positive opinions present. 0 otherwise
        '''
        count = 0
        positive = 0
        positive_ops = [x.lower() for x in self.positive_opinions]
        for w in tokens:
            if w.lower() in positive_ops:
                count += 1
        if count > 0:
            positive = 1
        return count, positive

    def hasVulgarWords(self,tokens):
        '''
        Checks if vulgar words from online list present
        :return: 1 if present
        '''
        for w in tokens:
            if w.lower() in self.vulgar_words:
                return 1
        return 0

    def hasAbbreviations(self,tokens):
        '''
        Checks if abbreviations from online list present
        :return: 1 if present
        '''
        for w in tokens:
            if w.lower() in self.web_abbreviations:
                return 1
        return 0

    def hasTwitterJargons(self,tokens):
        '''
        Checks if twitter specific abbreviations and jargons from online list present
        :return: 1 if present
        '''
        for w in tokens:
            if w in self.twitter_jargons:
                return 1
        return 0

    def hasEmoticons(self,tokens):
        '''
        Checks if abbreviations from online list present
        :return: 1 if present
        '''
        for w in tokens:
            if w in self.emoticons_list:
                return 1
        return 0

    def isInterrogative(self,tokens):
        '''
        Checks if ? present
        :return: 1 if present
        '''
        for term in tokens:
            if term in string.punctuation:
                if term == "?":
                    return 1
        return 0

    def isExclamatory(self,tokens):
        '''
        Checks if ! present
        :return: 1 if present
        '''
        for term in tokens:
            if term in string.punctuation:
                if term == "!":
                    return 1
        return 0

    def hasHash(self,tokens):
        '''
        Checks if # present
        :return: 1 if present, 0 otherwise; 1 if present in the beginning, 0 oherwise
        '''
        for index, term in enumerate(tokens):
            if re.match("#\w*", term):
                if 2 * index < len(tokens):
                    return 1,1
                else:
                    return 1,0
        return 0,0

    def hasRT(self,tokens):
        '''
        Checks if RT present
        :return: 1 if present, 0 otherwise; 1 if present in the beginning, 0 oherwise
        '''
        for index, term in enumerate(tokens):
           if term == "RT":
                if 2 * index < len(tokens):
                    return 1,1
                else:
                    return 1,0
        return 0,0

    def hasATag(self,tokens):
        '''
        Checks if @ present
        :return: 1 if present, 0 otherwise; 1 if present in the beginning, 0 oherwise
        '''
        for index, term in enumerate(tokens):
            if re.match("@\w*", term):
                if 2 * index < len(tokens):
                    return 1,1
                else:
                    return 1,0
        return 0,0

    def hasALink(self,tokens):
        """
        Checks if the tweet has a url
        :return: 1 if url present
        """
        for w in tokens:
            if re.findall('http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\(\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+', w):
                return 1
        return 0

    def hasManyNumbers(self,tokens):
        """
        Checks if the tweet many numbers
        :return: 1 if more than 5 numbers present
        """
        count = 0
        for w in tokens:
            if re.match('[\d]+',w):
               count = count + 1
        if count > 2:
            return 1
        else:
            return 0

    def hasManyNonAscii(self, tokens):
        """
        Checks if the tweet has non-ascii characters
        :return: 1 if more than 5 such present
        """
        count = 0
        for token in tokens:
            if re.match('[^\x00-\x7F]',token):
                count = count + 1
        if count > 5:
            return 1
        else:
            return 0

    def hasLinksToReputedDomains(self, url_string):
        """
        Checks if the tweet links to some trusted pages
        :return: 1 if more any such link present
        """
        for i in self.trusted_domains:
            pattern = re.compile(i+'.*')
            if re.findall(pattern,url_string):
                return 1
        return 0

    def hasSpeechActVerbs(self,tokens):
        '''
        Returns 265 binary features for the speech act verbs
        :return: a dictionary
        '''
        speech_act_feature_dict = {}        # dict stores the 265 features
        for w in self.speech_act_verbs:     # initializing
            speech_act_feature_dict[w] = 0
        pos_list = nltk.pos_tag(tokens)
        for entry in pos_list:
            if entry[1] in self.verb_tags:   # check if the tokens that are verbs are speech act verbs
                verb = self.stemmer.stem(entry[0].lower())
                if verb in self.speech_act_verbs:
                    speech_act_feature_dict[verb] = 1
        return speech_act_feature_dict

    def collectNGramFeatures(self,tweet):
        '''
        Collects the n-grams present in tweet that also occur through out the document more than 5 times
        :return: a count matrix
        '''

        # initializing matrix to all 0s
        tweet_n_gram_count_dict = {}

        for gram in self.n_gram_count_matrix:
            tweet_n_gram_count_dict[gram] = 0
            p = re.compile(re.escape(gram))
            if re.findall(p, str(tweet)):
                tweet_n_gram_count_dict[gram] = 1

        return tweet_n_gram_count_dict

    def returnUnigramMatrix(self,collection):
        '''
        Reurns unigrams in document- done to fetch Brexit related keywords
        '''
        X = self.vectorizer_unigram.fit_transform(collection)
        feature_names = self.vectorizer_unigram.get_feature_names()
        term_freqs = X.sum(axis=0).A1
        unigram_count_matrix = dict(zip(feature_names, term_freqs))

        return unigram_count_matrix

    def calculateNGramsEntropy(self, feature_names, SA_tagged_collection ):
        '''
        Calculates entropy of n-grams (uni, bi and tri grams)
        :param feature_names: n-grams
        :param SA_tagged_collection: speect act tagged collection
        '''

        #Remove n-grams with #,@,RT, only numbers
        for i in feature_names:
            if re.findall('#.*', i):
                feature_names.remove(i)
            elif re.findall('[\d]+', i):
                feature_names.remove(i)
            elif re.findall('@.*', i):
                feature_names.remove(i)
            elif re.findall('RT.*', i):
                feature_names.remove(i)
            elif re.findall('http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\(\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+',i):
                feature_names.remove(i)

        # Remove potential Brexit keywords
        for i in feature_names:
            for j in i.split():
                if j in self.brexit_keywords:
                    feature_names.remove(i)
                    break

        for term in nltk.pos_tag(feature_names):
            # Remove proper nouns
            if term[1] == "NNP" or term[1] == "NNPS" or term[1] == "NN":
                feature_names.remove(term[0])

        #Find entropy of each n-gram
        entropy = {}
        for item in feature_names:
            entropy[item] = 0
        total = len(SA_tagged_collection)
        occurrences = {}
        tagged_tweets= {}
        for gram in SA_tagged_collection:
            if SA_tagged_collection.get(gram)[1] != '':
                tagged_tweets[gram] = SA_tagged_collection.get(gram)

        for item in feature_names:
            p = re.compile(re.escape(item))
            for gram in tagged_tweets:
                occurrences[tagged_tweets.get(gram)[1]] = 0
            for gram in tagged_tweets:
                times_in_tweet = len(re.findall(p, str(tagged_tweets.get(gram)[0])))
                occurrences[tagged_tweets.get(gram)[1]] += times_in_tweet
            for gram in tagged_tweets:
                entropy[item] += -(occurrences[tagged_tweets.get(gram)[1]]/total)*math.log((occurrences[tagged_tweets.get(gram)[1]]+1)/total)

        # n-gram entropies
        with open('../Data/test/entropy_raw.csv', 'w+', newline='', encoding='utf8') as out_file:
            writer = csv.writer(out_file)
            for w in sorted(entropy, key=entropy.get):
                writer.writerow([w, entropy.get(w)])

        # entropies greater than 0 and less than 0.02 chosen.
        with open('../Data/test/N_grams_final.csv', 'w+', newline='', encoding='utf8') as out_file:
            writer = csv.writer(out_file)
            for w in entropy:
                if entropy.get(w) > 0 and entropy.get(w) < 0.02 :
                    writer.writerow([w])

    def createNGramCountMatrix(self,collection,SA_tagged_collection, create_n_grams):
        '''
        Creates a n-gram count matrix of all the tweets. N-grams that have no common nouns and occur more than 5 times
        :return: a count matrix
        '''
        X = self.vectorizer.fit_transform(collection)
        feature_names = self.vectorizer.get_feature_names()
        print("Size of n grams =  %d " % len(feature_names))

        if create_n_grams :
            self.calculateNGramsEntropy(feature_names, SA_tagged_collection)
        else:
            # USING good n-grams generated beforehand
            feature_names = [line.rstrip('\n') for line in open('../Data/Lists/N_grams_final.csv', encoding='utf8')]

        term_freqs = X.sum(axis=0).A1
        self.n_gram_count_matrix = dict(zip(feature_names, term_freqs))

        # for w in self.n_gram_count_matrix.copy():
        #   #keep only those n-grams that have a frequency > 5
        #     if(self.n_gram_count_matrix.get(w) < 5):
        #         self.n_gram_count_matrix.pop(w)

        #
        # # Normalize by dividing by log of number of occurences on n-gram in collection
        # for gram in entropy.copy():
        #     if gram in self.n_gram_count_matrix.keys():
        #         entropy[gram] = entropy[gram]/(math.log(self.n_gram_count_matrix.get(gram)))
        #     else:
        #         entropy.pop(gram)
        #

        # # n-gram entropies
        # with open('../Data/test/entropy.csv', 'w+', newline='', encoding='utf8') as out_file:
        #     writer = csv.writer(out_file)
        #     for w in sorted(entropy, key=entropy.get):
        #         writer.writerow([w, entropy.get(w)])

        return self.tweetFeatures()

    def speechActTagTweet(self,tweet):
        tokens = self.tokenize(tweet)
        pos_list = nltk.pos_tag(tokens)
        tag = ''
        for entry in pos_list:
            if entry[1] in self.verb_tags:  # check if the tokens that are verbs are speech act verbs
                verb = self.stemmer.stem(entry[0].lower())
                if verb in self.speech_act_verbs:
                    tag = verb
                    break
        return tag

    def speechActTagCollection(self, collection):
        self.speech_act_tags = {}
        for tweet in collection:
            self.speech_act_tags[tweet] = (collection.get(tweet), self.speechActTagTweet(collection.get(tweet)))
        return self.speech_act_tags

    def enrichTweets(self, tweet, urls):
        """
        Features added to a row of data.
        :return: Returns a row with features added
        """
        row=[]
        basic_row = []
        binary_ngrams_row= []
        #tokenize tweet
        tokens = self.tokenize(tweet)
        # remove stop words
        tokens = self.removeStopWords(tokens)
        # add  Vulgar words feature
        row.append(self.hasVulgarWords(tokens))
        binary_ngrams_row.append(self.hasVulgarWords(tokens))
        basic_row.append(self.hasVulgarWords(tokens))
        # add Emoticons feature
        row.append(self.hasEmoticons(tokens))
        basic_row.append(self.hasEmoticons(tokens))
        binary_ngrams_row.append(self.hasEmoticons(tokens))
        # add Interrogation feature
        row.append(self.isInterrogative(tokens))
        binary_ngrams_row.append(self.isInterrogative(tokens))
        basic_row.append(self.isInterrogative(tokens))
        # add Exclamation feature
        row.append((self.isExclamatory(tokens)))
        binary_ngrams_row.append((self.isExclamatory(tokens)))
        basic_row.append((self.isExclamatory(tokens)))
        # add Abbreviations feature
        row.append(self.hasAbbreviations(tokens))
        basic_row.append(self.hasAbbreviations(tokens))
        binary_ngrams_row.append(self.hasAbbreviations(tokens))
        # add twitter jargons feature
        row.append(self.hasTwitterJargons(tokens))
        basic_row.append(self.hasTwitterJargons(tokens))
        binary_ngrams_row.append(self.hasTwitterJargons(tokens))
        # add Twiiter specific characters' features- presence and position
        presence, position = self.hasHash(tokens)
        row.append(presence)
        row.append(position)

        basic_row.append(presence)
        basic_row.append(position)

        binary_ngrams_row.append(presence)
        binary_ngrams_row.append(position)

        presence, position = self.hasATag(tokens)
        row.append(presence)
        row.append(position)

        basic_row.append(presence)
        basic_row.append(position)

        binary_ngrams_row.append(presence)
        binary_ngrams_row.append(position)

        presence, position = self.hasRT(tokens)
        row.append(presence)
        row.append(position)

        basic_row.append(presence)
        basic_row.append(position)

        binary_ngrams_row.append(presence)
        binary_ngrams_row.append(position)

        # add URL presence feature
        row.append(self.hasALink(tokens))
        basic_row.append(self.hasALink(tokens))
        binary_ngrams_row.append(self.hasALink(tokens))
        #add speech act verbs features
        sac_feature_dict = self.hasSpeechActVerbs(tokens)
        for verb in self.speech_act_verbs:
              row.append(sac_feature_dict.get(verb))
              basic_row.append(sac_feature_dict.get(verb))
              binary_ngrams_row.append(sac_feature_dict.get(verb))

        #more than 5 numbers in tweet
        row.append(self.hasManyNumbers(tokens))
        basic_row.append(self.hasManyNumbers(tokens))
        binary_ngrams_row.append(self.hasManyNumbers(tokens))

        #presence of non-ascii characters
        row.append(self.hasManyNonAscii(tokens))
        basic_row.append(self.hasManyNonAscii(tokens))
        binary_ngrams_row.append(self.hasManyNonAscii(tokens))

        #presence of links to trusted domains
        row.append(self.hasLinksToReputedDomains(urls))
        binary_ngrams_row.append(self.hasLinksToReputedDomains(urls))
        basic_row.append(self.hasLinksToReputedDomains(urls))

        # add positive and negative opinion feature
        pos_count, pos_presence = self.hasNegativeOpinions(tokens)
        neg_count, neg_presence = self.hasPositiveOpinions(tokens)

        # counts
        row.append(pos_count)
        row.append(neg_count)

        # boolean- positive/negative opinion feature
        basic_row.append(pos_presence)
        basic_row.append(neg_presence)

        binary_ngrams_row.append(pos_presence)
        binary_ngrams_row.append(neg_presence)

        #sentiment feature
        pos_per, neg_per = self.sentiment(tweet)
        tweet_sentiment = 1
        if pos_per > neg_per:
            tweet_sentiment = 1
        elif pos_per < neg_per:
            tweet_sentiment = 0

        row.append(tweet_sentiment)
        basic_row.append(tweet_sentiment)

        # n grams from tweet
        n_gram_feature_dict = self.collectNGramFeatures(tweet)
        for w in n_gram_feature_dict:
            row.append(n_gram_feature_dict.get(w))
            binary_ngrams_row.append(n_gram_feature_dict.get(w))

        return row, basic_row, binary_ngrams_row

    def enrichTweetsWithNgrams(self, tweet):
        '''
        Returns only n gram features for tweet
        '''
        row = []
        n_gram_feature_dict = self.collectNGramFeatures(tweet)
        for w in n_gram_feature_dict:
            row.append(n_gram_feature_dict.get(w))
        return row

    def hasPositiveEmoticons(self, tokens):
        '''
        Checks if abbreviations from online list present
        :return: 1 if present
        '''
        for w in tokens:
            if w in self.pos_emoticons_list:
                return 1
        return 0

    def hasNegativeEmoticons(self, tokens):
        '''
        Checks if abbreviations from online list present
        :return: 1 if present
        '''
        for w in tokens:
            if w in self.neg_emoticons_list:
                return 1
        return 0

    def sentiment(self, text):
        length = len(text)
        tokens = self.tokenize(text)

        tokens = [w.lower() for w in tokens if not w in set(stopwords.words('english') + self.web_abbreviations + list(string.punctuation))]
        tokens = list(set(tokens))
        positive_count, negative_count = self.simin_opinion_count(tokens)
        return positive_count, negative_count

    def simin_opinion_count(self, tokens):
        '''
                Checks if positive opnions present
                :return: 1 if positive opinions present. 0 otherwise
                '''
        pos_count = 0
        neg_count = 0

        for w in tokens:
            synonyms = []
            for syn in nltk.corpus.wordnet.synsets(w):
                for l in syn.lemmas():
                   synonyms.append(l.name())
            for syn in synonyms:
                if syn in self.positive_ops:
                    pos_count += 1
                    break
                elif syn in self.negative_ops:
                    neg_count += 1
                    break
            if w in self.neg_emoticons_list:
                neg_count += 1
            if w in self.pos_emoticons_list:
                pos_count += 1
        return pos_count, neg_count

    def hasFirstPersonPronouns(self, tokens):
        for w in tokens:
            if w in self.first_person_pronouns:
                return 1
        return 0


import discord

from discord import slash_command
from discord.ext import commands

from modals.Anonymous import Anonymous
from utils import globals as GG

log = GG.log

class AnonModal(commands.Cog):
    def __init__(self, bot):
        self.bot = bot

    @slash_command(name="anonreport")
    @commands.guild_only()
    async def anonreport(self, ctx):
        """Opens a modal to post an anonymous report."""
        delivery_channel = await GG.MDB['channelinfo'].find_one({"guild": ctx.interaction.guild_id, "type": "DELIVERY"})
        if delivery_channel is None:
            await self.noDeliveryChannel(ctx.interaction.guild, ctx.interaction)
            await self.notifyOwner(ctx.interaction.guild, ctx.interaction)
            return

        thread_channel = await GG.MDB['channelinfo'].find_one({"guild": ctx.interaction.guild_id, "type": "ANON"})
        if thread_channel is None:
            await self.noThreadChannel(ctx.interaction.guild, ctx.interaction)
            await self.notifyOwner(ctx.interaction.guild, ctx.interaction)
            return

        delivery_channel = await self.bot.fetch_channel(delivery_channel['channel'])
        thread_channel = await self.bot.fetch_channel(thread_channel['channel'])
        modal = Anonymous(self.bot, ctx.interaction, ctx.interaction.user, delivery_channel, thread_channel)
        await ctx.interaction.response.send_modal(modal)

    async def notifyOwner(self, guild, interaction_for_modal=None):
        ServerOwner = self.bot.get_user(guild.owner.id)
        if ServerOwner.dm_channel is None:
            dm_channel = await ServerOwner.create_dm()
        else:
            dm_channel = ServerOwner.dm_channel
        Embed = discord.Embed(title=f"No delivery channel and/or thread channel exists.",
                              description=f"Hello, a user of your server {guild.name} tried to make an "
                                          f"anonymous report, but you have yet to make a delivery channel and/or thread channel. You can do this by "
                                          f"executing the following command: ``/addchannel <id> DELIVERY`` and/or ``/addchannel <id> ANON``. Thank you!")
        try:
            await dm_channel.send(embed=Embed)
        except discord.Forbidden:
            if interaction_for_modal is not None:
                await interaction_for_modal.channel.send(
                    f"Hello {guild.owner.mention}, a user just tried to make an anonymous report, but you have "
                    f"yet to make a delivery channel. You can do this by executing the following command: ``/addchannel "
                    f"<id> DELIVERY``. Thank you!")

    async def noThreadChannel(self, guild, interaction_for_modal=None):
        Reporter = self.bot.get_user(interaction_for_modal.user.id)
        if Reporter.dm_channel is None:
            dm_channel = await Reporter.create_dm()
        else:
            dm_channel = Reporter.dm_channel
        Embed = discord.Embed(title=f"Error",
                              description="Sorry this Server doesnt have a thread channel yet, "
                                          "I will notify the Server Owner for you if I'm able.")
        try:
            await dm_channel.send(embed=Embed)
        except discord.Forbidden:
            if interaction_for_modal is not None:
                await interaction_for_modal.channel.send(
                    "@Reporter, you have turned off the ability to let me DM you. I will notify the "
                    "Server Owner that they haven't created a delivery channel yet.")

    async def noDeliveryChannel(self, guild, interaction_for_modal=None):
        Reporter = self.bot.get_user(interaction_for_modal.user.id)
        if Reporter.dm_channel is None:
            dm_channel = await Reporter.create_dm()
        else:
            dm_channel = Reporter.dm_channel
        Embed = discord.Embed(title=f"Error",
                              description="Sorry this Server doesnt have a delivery channel yet, "
                                          "I will notify the Server Owner for you if I'm able.")
        try:
            await dm_channel.send(embed=Embed)
        except discord.Forbidden:
            if interaction_for_modal is not None:
                await interaction_for_modal.channel.send(
                    "@Reporter, you have turned off the ability to let me DM you. I will notify the "
                    "Server Owner that they haven't created a delivery channel yet.")

def setup(bot):
    log.info("[Cog] AnonModal Report")
    bot.add_cog(AnonModal(bot))


from typing import List, Union
import numpy as np
from sentence_transformers import SentenceTransformer
from services.components.cos_similarity import cos_sim
from services.interface.abstract_model import PlagiarismDriver

class Similarity(PlagiarismDriver):
    def __init__(
            self,
            model_name_or_path="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
            device=None):
        if isinstance(model_name_or_path, str):
            self.plagiarism_model = SentenceTransformer(
                model_name_or_path,
                device=device
            )
        elif hasattr(model_name_or_path, "encode"):
            self.plagiarism_model = model_name_or_path
        else:
            raise ValueError("model_name_or_path is transformers model name or path")
    def _get_vector(
            self,
            sentences: Union[str, List[str]],
            batch_size: int = 64,
            show_progress_bar: bool = False,
    ) -> np.ndarray:
        """
        Returns the embeddings for a batch of sentences
        """
        return self.plagiarism_model.encode(sentences, batch_size=batch_size, show_progress_bar=show_progress_bar)

    def similarity(self, a: Union[str, List[str]], b: Union[str, List[str]]):
        """
        Compute similarity between two texts.
        """
        text_emb1 = self._get_vector(a)
        text_emb2 = self._get_vector(b)
        return cos_sim(text_emb1, text_emb2)


import math
import torch
from torch import nn

from vqfr.archs.stylegan_arch import UpFirDnSmooth
from vqfr.ops.upfirdn2d import upfirdn2d
from vqfr.utils.registry import ARCH_REGISTRY

def get_haar_wavelet(in_channels):
    haar_wav_l = 1 / (2**0.5) * torch.ones(1, 2)
    haar_wav_h = 1 / (2**0.5) * torch.ones(1, 2)
    haar_wav_h[0, 0] = -1 * haar_wav_h[0, 0]

    haar_wav_ll = haar_wav_l.T * haar_wav_l
    haar_wav_lh = haar_wav_h.T * haar_wav_l
    haar_wav_hl = haar_wav_l.T * haar_wav_h
    haar_wav_hh = haar_wav_h.T * haar_wav_h

    return haar_wav_ll, haar_wav_lh, haar_wav_hl, haar_wav_hh

def dwt_init(x):
    x01 = x[:, :, 0::2, :] / 2
    x02 = x[:, :, 1::2, :] / 2
    x1 = x01[:, :, :, 0::2]
    x2 = x02[:, :, :, 0::2]
    x3 = x01[:, :, :, 1::2]
    x4 = x02[:, :, :, 1::2]
    x_LL = x1 + x2 + x3 + x4
    x_HL = -x1 - x2 + x3 + x4
    x_LH = -x1 + x2 - x3 + x4
    x_HH = x1 - x2 - x3 + x4

    return torch.cat((x_LL, x_HL, x_LH, x_HH), 1)

def iwt_init(x):
    r = 2
    in_batch, in_channel, in_height, in_width = x.size()
    # print([in_batch, in_channel, in_height, in_width])
    out_batch, out_channel, out_height, out_width = (
        in_batch,
        int(in_channel / (r**2)),
        r * in_height,
        r * in_width,
    )
    x1 = x[:, 0:out_channel, :, :] / 2
    x2 = x[:, out_channel:out_channel * 2, :, :] / 2
    x3 = x[:, out_channel * 2:out_channel * 3, :, :] / 2
    x4 = x[:, out_channel * 3:out_channel * 4, :, :] / 2

    h = torch.zeros([out_batch, out_channel, out_height, out_width]).float().cuda()

    h[:, :, 0::2, 0::2] = x1 - x2 - x3 + x4
    h[:, :, 1::2, 0::2] = x1 - x2 + x3 - x4
    h[:, :, 0::2, 1::2] = x1 + x2 - x3 - x4
    h[:, :, 1::2, 1::2] = x1 + x2 + x3 + x4

    return h

class HaarTransform(nn.Module):

    def __init__(self, in_channels):
        super().__init__()

        ll, lh, hl, hh = get_haar_wavelet(in_channels)

        self.register_buffer('ll', ll)
        self.register_buffer('lh', lh)
        self.register_buffer('hl', hl)
        self.register_buffer('hh', hh)

    def forward(self, input):
        ll = upfirdn2d(input, self.ll, down=2)
        lh = upfirdn2d(input, self.lh, down=2)
        hl = upfirdn2d(input, self.hl, down=2)
        hh = upfirdn2d(input, self.hh, down=2)

        return torch.cat((ll, lh, hl, hh), 1)

class InverseHaarTransform(nn.Module):

    def __init__(self, in_channels):
        super().__init__()

        ll, lh, hl, hh = get_haar_wavelet(in_channels)

        self.register_buffer('ll', ll)
        self.register_buffer('lh', -lh)
        self.register_buffer('hl', -hl)
        self.register_buffer('hh', hh)

    def forward(self, input):
        ll, lh, hl, hh = input.chunk(4, 1)
        ll = upfirdn2d(ll, self.ll, up=2, pad=(1, 0, 1, 0))
        lh = upfirdn2d(lh, self.lh, up=2, pad=(1, 0, 1, 0))
        hl = upfirdn2d(hl, self.hl, up=2, pad=(1, 0, 1, 0))
        hh = upfirdn2d(hh, self.hh, up=2, pad=(1, 0, 1, 0))

        return ll + lh + hl + hh

class ConvLayer(nn.Sequential):
    """Conv Layer used in StyleGAN2 Discriminator.
    Args:
        in_channels (int): Channel number of the input.
        out_channels (int): Channel number of the output.
        kernel_size (int): Kernel size.
        downsample (bool): Whether downsample by a factor of 2.
            Default: False.
        resample_kernel (list[int]): A list indicating the 1D resample
            kernel magnitude. A cross production will be applied to
            extent 1D resample kernel to 2D resample kernel.
            Default: (1, 3, 3, 1).
        bias (bool): Whether with bias. Default: True.
        activate (bool): Whether use activateion. Default: True.
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size,
                 downsample=False,
                 resample_kernel=(1, 3, 3, 1),
                 bias=True,
                 activate=True,
                 equalize=True):
        layers = []
        # downsample
        if downsample:
            layers.append(
                UpFirDnSmooth(resample_kernel, upsample_factor=1, downsample_factor=2, kernel_size=kernel_size))
            stride = 2
            self.padding = 0
        else:
            stride = 1
            self.padding = kernel_size // 2
        # conv
        if equalize:
            raise NotImplementedError
        else:
            layers.append(
                SpectralNorm(
                    nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, padding=self.padding, bias=bias)))
            # activation
            if activate:
                layers.append(nn.LeakyReLU(0.2))

        super(ConvLayer, self).__init__(*layers)

class ConvBlock(nn.Module):

    def __init__(self, in_channel, out_channel, blur_kernel=[1, 3, 3, 1], equalize=True):
        super().__init__()

        self.conv1 = ConvLayer(in_channel, in_channel, 3, equalize=equalize)
        self.conv2 = ConvLayer(in_channel, out_channel, 3, downsample=True, equalize=equalize)

    def forward(self, input):
        out = self.conv1(input)
        out = self.conv2(out)

        return out

def make_kernel(k):
    k = torch.tensor(k, dtype=torch.float32)

    if k.ndim == 1:
        k = k[None, :] * k[:, None]

    k /= k.sum()

    return k

class Downsample(nn.Module):

    def __init__(self, kernel, factor=2):
        super().__init__()

        self.factor = factor
        kernel = make_kernel(kernel)
        self.register_buffer('kernel', kernel)

        p = kernel.shape[0] - factor

        pad0 = (p + 1) // 2
        pad1 = p // 2

        self.pad = (pad0, pad1)

    def forward(self, input):
        out = upfirdn2d(input, self.kernel, up=1, down=self.factor, pad=self.pad)

        return out

class FromRGB(nn.Module):

    def __init__(self, out_channel, downsample=True, blur_kernel=[1, 3, 3, 1], equalize=True):
        super().__init__()

        self.downsample = downsample

        if downsample:
            self.iwt = InverseHaarTransform(3)
            self.downsample = Downsample(blur_kernel)
            self.dwt = HaarTransform(3)

        self.conv = ConvLayer(3 * 4, out_channel, 3, equalize=equalize)

    def forward(self, input, skip=None):
        if self.downsample:
            input = self.iwt(input)
            input = self.downsample(input)
            input = self.dwt(input)

        out = self.conv(input)

        if skip is not None:
            out = out + skip

        return input, out

def l2normalize(v, eps=1e-12):
    return v / (v.norm() + eps)

class SpectralNorm(nn.Module):

    def __init__(self, module, name='weight', power_iterations=1):
        super(SpectralNorm, self).__init__()
        self.module = module
        self.name = name
        self.power_iterations = power_iterations
        if not self._made_params():
            self._make_params()

    def _update_u_v(self):
        u = getattr(self.module, self.name + '_u')
        v = getattr(self.module, self.name + '_v')
        w = getattr(self.module, self.name + '_bar')

        height = w.data.shape[0]
        for _ in range(self.power_iterations):
            v.data = l2normalize(torch.mv(torch.t(w.view(height, -1).data), u.data))
            u.data = l2normalize(torch.mv(w.view(height, -1).data, v.data))

        # sigma = torch.dot(u.data, torch.mv(w.view(height,-1).data, v.data))
        sigma = u.dot(w.view(height, -1).mv(v))
        setattr(self.module, self.name, w / sigma.expand_as(w))

    def _made_params(self):
        try:
            getattr(self.module, self.name + '_u')
            getattr(self.module, self.name + '_v')
            getattr(self.module, self.name + '_bar')
            return True
        except AttributeError:
            return False

    def _make_params(self):
        w = getattr(self.module, self.name)

        height = w.data.shape[0]
        width = w.view(height, -1).data.shape[1]

        u = nn.Parameter(w.data.new(height).normal_(0, 1), requires_grad=False)
        v = nn.Parameter(w.data.new(width).normal_(0, 1), requires_grad=False)
        u.data = l2normalize(u.data)
        v.data = l2normalize(v.data)
        w_bar = nn.Parameter(w.data)

        del self.module._parameters[self.name]

        self.module.register_parameter(self.name + '_u', u)
        self.module.register_parameter(self.name + '_v', v)
        self.module.register_parameter(self.name + '_bar', w_bar)

    def forward(self, *args):
        self._update_u_v()
        return self.module.forward(*args)

@ARCH_REGISTRY.register()
class SWADiscriminator(nn.Module):

    def __init__(self, out_size, channel_multiplier=2, resample_kernel=[1, 3, 3, 1]):
        super().__init__()

        channels = {
            4: 512,
            8: 512,
            16: 512,
            32: 512,
            64: 256 * channel_multiplier,
            128: 128 * channel_multiplier,
            256: 64 * channel_multiplier,
            512: 32 * channel_multiplier,
            1024: 16 * channel_multiplier,
        }

        self.dwt = HaarTransform(3)

        self.from_rgbs = nn.ModuleList()
        self.convs = nn.ModuleList()

        log_size = int(math.log(out_size, 2)) - 1

        in_channel = channels[out_size]

        for i in range(log_size, 2, -1):
            out_channel = channels[2**(i - 1)]

            self.from_rgbs.append(FromRGB(in_channel, downsample=i != log_size, equalize=False))
            self.convs.append(ConvBlock(in_channel, out_channel, resample_kernel, equalize=False))

            in_channel = out_channel

        self.from_rgbs.append(FromRGB(channels[4], equalize=False))

        self.stddev_group = 4
        self.stddev_feat = 1

        self.final_conv = ConvLayer(in_channel + 1, channels[4], 3, equalize=False)

        self.final_linear = nn.Sequential(
            SpectralNorm(nn.Linear(channels[4] * 4 * 4, channels[4])),
            nn.LeakyReLU(0.2),
            SpectralNorm(nn.Linear(channels[4], 1)),
        )

    def forward(self, input):
        input = self.dwt(input)
        out = None

        for from_rgb, conv in zip(self.from_rgbs, self.convs):
            input, out = from_rgb(input, out)
            out = conv(out)

        _, out = self.from_rgbs[-1](input, out)

        batch, channel, height, width = out.shape
        group = min(batch, self.stddev_group)
        stddev = out.view(group, -1, self.stddev_feat, channel // self.stddev_feat, height, width)
        stddev = torch.sqrt(stddev.var(0, unbiased=False) + 1e-8)
        stddev = stddev.mean([2, 3, 4], keepdims=True).squeeze(2)
        stddev = stddev.repeat(group, 1, height, width)
        out = torch.cat([out, stddev], 1)

        out = self.final_conv(out)

        out = out.view(batch, -1)
        out = self.final_linear(out)

        return out