import os
import time
import copy
import random
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

import cv2 as cv
import PIL
from PIL import Image
import PIL.Image as pil_image
# import collections
# import os.path as osp
# import scipy
# import scipy.misc
# import imageio
import h5py
from tqdm import tqdm

import torch
import torch.nn as nn11
import torch.nn.functional as F
import torch.optim as optim
import torch.backends.cudnn as cudnn
from torch.utils.data import Dataset
from torch.utils.data.dataloader import DataLoader

import torchvision

# path = r'D:\Python\★★Python_POSTECH_AI\Postech_AI 7) Computer_Vision\Dataset'
# path = r'D:/Python/★★Python_POSTECH_AI/Postech_AI 7) Computer_Vision/Dataset/Lecture08_Image restoration/'
path = r'/home/pirl/data/Lecture08_Image restoration/'
origin_path = os.getcwd()
os.chdir(path)
# os.listdir()

seed = 1
no_cuda = False
torch.manual_seed(seed)
use_cuda = not no_cuda and torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
kwargs = {'num_workers': 1, 'pin_memory':True} if use_cuda else {}
torch.cuda.empty_cache()
# print(torch.cuda.memory_summary())
# torch.cuda.empty_cache()    # 메모리 초기화

# Train Dataset 불러오는 class
class TrainDataset(Dataset):
    def __init__(self, h5_file):
        super(TrainDataset, self).__init__()
        self.h5_file = h5_file

    def __getitem__(self, idx):
        with h5py.File(self.h5_file, 'r') as f:
            return np.expand_dims(f['lr'][idx] / 255., 0), np.expand_dims(f['hr'][idx] / 255., 0)

    def __len__(self):
        with h5py.File(self.h5_file, 'r') as f:
            return len(f['lr'])

# Evaluation Dataset 불러오는 class
class EvalDataset(Dataset):
    def __init__(self, h5_file):
        super(EvalDataset, self).__init__()
        self.h5_file = h5_file

    def __getitem__(self, idx):
        with h5py.File(self.h5_file, 'r') as f:
            return np.expand_dims(f['lr'][str(idx)][:, :] / 255., 0), np.expand_dims(f['hr'][str(idx)][:, :] / 255., 0)

    def __len__(self):
        with h5py.File(self.h5_file, 'r') as f:
            return len(f['lr'])

# # Model definition
#
# ## TO DO - Define the model by following the given architecture
# ### SRCNN
# conv1: (in_channel) num_channels, (out_channel) 64, (kernel_size) 9, (stride) 1, (padding) 4
# conv2: (in_channel) 64, (out_channel) 32, (kernel_size) 5, (stride) 1, (padding) 2
# conv3: (in_channel) 32, (out_channel) num_channels, (kernel_size) 5, (stride) 1, (padding) 2
#
# Each layers should have ReLU activation but the last conv layer should not have any activation layer.

# SRCNN **** ============================================================================
class SRCNN(nn.Module):
    def __init__(self, num_channels=1):
        super(SRCNN, self).__init__()
        self.conv1 = nn.Conv2d(num_channels, 64, 9, 1, 4)
        self.conv2 = nn.Conv2d(64, 32, 5, 1, 2)
        self.conv3 = nn.Conv2d(32, num_channels, 5, 1, 2)

    def forward(self, x):
        x = F.relu(self.conv1(x), inplace=True)
        x = F.relu(self.conv2(x), inplace=True)
        x = self.conv3(x)
        return x

# # Util functions
def convert_rgb_to_y(img):
    if type(img) == np.ndarray:
        return 16. + (64.738 * img[:, :, 0] + 129.057 * img[:, :, 1] + 25.064 * img[:, :, 2]) / 256.
    elif type(img) == torch.Tensor:
        if len(img.shape) == 4:
            img = img.squeeze(0)
        return 16. + (64.738 * img[0, :, :] + 129.057 * img[1, :, :] + 25.064 * img[2, :, :]) / 256.
    else:
        raise Exception('Unknown Type', type(img))

def convert_rgb_to_ycbcr(img):
    if type(img) == np.ndarray:
        y = 16. + (64.738 * img[:, :, 0] + 129.057 * img[:, :, 1] + 25.064 * img[:, :, 2]) / 256.
        cb = 128. + (-37.945 * img[:, :, 0] - 74.494 * img[:, :, 1] + 112.439 * img[:, :, 2]) / 256.
        cr = 128. + (112.439 * img[:, :, 0] - 94.154 * img[:, :, 1] - 18.285 * img[:, :, 2]) / 256.
        return np.array([y, cb, cr]).transpose([1, 2, 0])
    elif type(img) == torch.Tensor:
        if len(img.shape) == 4:
            img = img.squeeze(0)
        y = 16. + (64.738 * img[0, :, :] + 129.057 * img[1, :, :] + 25.064 * img[2, :, :]) / 256.
        cb = 128. + (-37.945 * img[0, :, :] - 74.494 * img[1, :, :] + 112.439 * img[2, :, :]) / 256.
        cr = 128. + (112.439 * img[0, :, :] - 94.154 * img[1, :, :] - 18.285 * img[2, :, :]) / 256.
        return torch.cat([y, cb, cr], 0).permute(1, 2, 0)
    else:
        raise Exception('Unknown Type', type(img))

def convert_ycbcr_to_rgb(img):
    if type(img) == np.ndarray:
        r = 298.082 * img[:, :, 0] / 256. + 408.583 * img[:, :, 2] / 256. - 222.921
        g = 298.082 * img[:, :, 0] / 256. - 100.291 * img[:, :, 1] / 256. - 208.120 * img[:, :, 2] / 256. + 135.576
        b = 298.082 * img[:, :, 0] / 256. + 516.412 * img[:, :, 1] / 256. - 276.836
        return np.array([r, g, b]).transpose([1, 2, 0])
    elif type(img) == torch.Tensor:
        if len(img.shape) == 4:
            img = img.squeeze(0)
        r = 298.082 * img[0, :, :] / 256. + 408.583 * img[2, :, :] / 256. - 222.921
        g = 298.082 * img[0, :, :] / 256. - 100.291 * img[1, :, :] / 256. - 208.120 * img[2, :, :] / 256. + 135.576
        b = 298.082 * img[0, :, :] / 256. + 516.412 * img[1, :, :] / 256. - 276.836
        return torch.cat([r, g, b], 0).permute(1, 2, 0)
    else:
        raise Exception('Unknown Type', type(img))

def calc_psnr(img1, img2):
    return 10. * torch.log10(1. / torch.mean((img1 - img2) ** 2))

class AverageMeter(object):
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

# # Training loop
train_file = "data/91-image_x3.h5"
eval_file = "data/Set5_x3.h5"
outputs_dir = "outputs/"

scale = 3
lr = 1e-4
batch_size = 16
num_workers = 8
seed = 123

outputs_dir = os.path.join(outputs_dir, 'x{}'.format(scale))

if not os.path.exists(outputs_dir):
    os.makedirs(outputs_dir)

cudnn.benchmark = True
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

torch.manual_seed(seed)

model = SRCNN().to(device)
criterion = nn.MSELoss()
optimizer = optim.Adam([
    {'params': model.conv1.parameters()},
    {'params': model.conv2.parameters()},
    {'params': model.conv3.parameters(), 'lr': lr * 0.1}
], lr=lr)

train_dataset = TrainDataset(train_file)

train_dataloader = DataLoader(dataset=train_dataset,
                                  batch_size=batch_size,
                                  shuffle=True,
                                  num_workers=num_workers,
                                  pin_memory=True,
                                  drop_last=True)
eval_dataset = EvalDataset(eval_file)
eval_dataloader = DataLoader(dataset=eval_dataset, batch_size=1)

# plt.imshow(a[0][1,0,:], 'gray')
# plt.imshow(a[1][1,0,:], 'gray')

# ======================================================================
# Input Size: 33 × 33 // 31 × 31
best_weights = copy.deepcopy(model.state_dict())
best_epoch = 0
best_psnr = 0.0

num_epochs = 10
for epoch in range(num_epochs):
    model.train()
    epoch_losses = AverageMeter()

    with tqdm(total=(len(train_dataset) - len(train_dataset) % batch_size)) as t:
        t.set_description('epoch: {}/{}'.format(epoch, num_epochs - 1))

        for data in train_dataloader:
            inputs, labels = data

            inputs = inputs.to(device)
            labels = labels.to(device)

            preds = model(inputs)

            loss = criterion(preds, labels)

            epoch_losses.update(loss.item(), len(inputs))

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

            t.set_postfix(loss='{:.6f}'.format(epoch_losses.avg))
            t.update(len(inputs))

    torch.save(model.state_dict(), os.path.join(outputs_dir, 'epoch_{}.pth'.format(epoch)))

    model.eval()
    epoch_psnr = AverageMeter()

    for data in eval_dataloader:
        inputs, labels = data

        inputs = inputs.to(device)
        labels = labels.to(device)

        with torch.no_grad():
            preds = model(inputs).clamp(0.0, 1.0)

        epoch_psnr.update(calc_psnr(preds, labels), len(inputs))

    print('eval psnr: {:.2f}'.format(epoch_psnr.avg))

    if epoch_psnr.avg > best_psnr:
        best_epoch = epoch
        best_psnr = epoch_psnr.avg
        best_weights = copy.deepcopy(model.state_dict())

print('best epoch: {}, psnr: {:.2f}'.format(best_epoch, best_psnr))
torch.save(best_weights, os.path.join(outputs_dir, 'best.pth'))

# SRGAN **** ============================================================================
# 1. SCRCNN, Image size에 상관없이 가능
#   Gaussian filter 적용하면 blur, image size와 무관
#   이와 비슷하게, SRCNN에서 구해진 conv filter는 super-resolution하는 filter

# 2. Perceptual Loss
# Idea : pixel값이 비슷하다고 우리 눈에 비슷하게 보이지만은 않는다
# 우리 눈에도 비슷하게 보이게 하도록 강제 하고 싶음.
# 우리눈 == Pretrained model (ImageNet Dataset에 train되어 있음)
# pretrained model의 앞부분 feature extractor를 사용하면 좋은 low-level feature를 얻을 수 있음
# pretrained model은 주로 VGG19를 많이 사용함

# feature를 가져와서 쓸때는 반드시 Relu Activation 다음에 있는 featuref를 쓰는게 좋음
# ... feature가 conv다음에 있을 수도있고, mapool 다음에도 있고, Relu 다음에도 있음

# 3. 왜 Pixelshuffle layer를 쓰는지? (deconv가 아니라)
# Deconv Layer를 쓸 경우 image generation task에서는 checkboard artifact가 생김, 이를 방지하기 위해 pixelshuffle layer를 사용 혹은 이후 upsampling

# 4. low Resolution dataset은 synthetic dataset
# High Resolution 이미지를 통해 bicubic downsampling을 통해 low Resolution image로 만들어줌

import glob
import random

from torch.autograd import Variable
import torchvision.transforms as transforms
from torchvision.utils import save_image, make_grid
from torchvision.models import vgg19

import argparse
import math
import itertools
import sys

# Normalization parameters for pre-trained PyTorch models
mean = np.array([0.485, 0.456, 0.406])
std = np.array([0.229, 0.224, 0.225])

# Image Dataset -----------------------------------------------------------------------------------
class ImageDataset(Dataset):
    def __init__(self, root, hr_shape):
        hr_height, hr_width = hr_shape
        # Transforms for low resolution images and high resolution images
        self.lr_transform = transforms.Compose(
            [
            transforms.Resize((hr_height//4, hr_width//4), Image.BICUBIC),  # Synthetic low resolution data를 만듦
            transforms.ToTensor(),
            transforms.Normalize(mean, std)
            ]
        )
        self.hr_transform = transforms.Compose(
            [
            transforms.Resize((hr_height, hr_width), Image.BICUBIC),
            transforms.ToTensor(),
            transforms.Normalize(mean, std)
            ]
        )

        self.files = sorted(glob.glob(root + "/*.*"))

    def __getitem__(self, index):
        img = Image.open(self.files[index % len(self.files)])
        img_lr = self.lr_transform(img)
        img_hr = self.hr_transform(img)

        return {"lr": img_lr, "hr": img_hr}

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

# Feature Extractor  -----------------------------------------------------------------------------------
class FeatureExtractor(nn.Module):
    '''
    input.shape: (batch, 3, width, hegiht)
    output.shape: (batch, 256, , width/4, hegiht/4)
    '''
    def __init__(self):
        super(FeatureExtractor, self).__init__()
        vgg19_model = vgg19(pretrained=True)
        self.feature_extractor = nn.Sequential(*list(vgg19_model.features.children())[:18])

    def forward(self, img):
        return self.feature_extractor(img)

# vgg19 = vgg19(pretrained=True)
# feature_extractor = nn.Sequential(*list(vgg19.features.children())[:18])

# Sequential(
#   (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#   (1): ReLU(inplace=True)
#   (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#   (3): ReLU(inplace=True)
#   (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#   (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#   (6): ReLU(inplace=True)
#   (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#   (8): ReLU(inplace=True)
#   (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#   (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#   (11): ReLU(inplace=True)
#   (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#   (13): ReLU(inplace=True)
#   (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#   (15): ReLU(inplace=True)
#   (16): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#   (17): ReLU(inplace=True)
# )
# a = FeatureExtractor()
# b = torch.rand((1,3,16,16))
# a(b).shape

class ResidualBlock(nn.Module):
    def __init__(self, in_features):
        super(ResidualBlock, self).__init__()
        self.conv_block = nn.Sequential(
            nn.Conv2d(in_features, in_features, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(in_features, 0.8),
            nn.PReLU(),
            nn.Conv2d(in_features, in_features, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(in_features, 0.8),
        )

    def forward(self, x):
        return x + self.conv_block(x)

# Generator ResNet -----------------------------------------------------------------------------------
class GeneratorResNet(nn.Module):
    '''
    input.shape: (batch, 3, width, hegiht)
    output.shape: (batch, 3, , width * 4, hegiht * 4)
    '''
    def __init__(self, in_channels=3, out_channels=3, n_residual_blocks=16):
        super(GeneratorResNet, self).__init__()

        # First layer
        self.conv1 = nn.Sequential(nn.Conv2d(in_channels, 64, kernel_size=9, stride=1, padding=4), nn.PReLU())

        # Residual blocks
        res_blocks = []
        for _ in range(n_residual_blocks):
            res_blocks.append(ResidualBlock(64))
        self.res_blocks = nn.Sequential(*res_blocks)

        # Second conv layer post residual blocks
        self.conv2 = nn.Sequential(nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1), nn.BatchNorm2d(64, 0.8))

        # Upsampling layers
        upsampling = []
        for out_features in range(2):
            upsampling += [
                # nn.Upsample(scale_factor=2),
                nn.Conv2d(64, 256, 3, 1, 1),
                nn.BatchNorm2d(256),
                nn.PixelShuffle(upscale_factor=2),      # Transpose Convolution과 동일한 역할 (PixelShuffle이 이미지를 만들때 더 안정적인)
                nn.PReLU(),
            ]
        self.upsampling = nn.Sequential(*upsampling)

        # Final output layer
        self.conv3 = nn.Sequential(nn.Conv2d(64, out_channels, kernel_size=9, stride=1, padding=4), nn.Tanh())

    def forward(self, x):
        out1 = self.conv1(x)
        out = self.res_blocks(out1)
        out2 = self.conv2(out)
        out = torch.add(out1, out2)
        out = self.upsampling(out)
        out = self.conv3(out)
        return out

# https://m.blog.naver.com/PostView.nhn?blogId=go21&logNo=221147971473&proxyReferer=https:%2F%2Fwww.google.com%2F
# nn.PixelShuffle(upscale_factor=3)     # upscale_factor = uf
# input.shape: (batch, channel, width, hegith)
#   → output.shape: (batch, channel/(uf**2), width*uf, height*uf)

# a = nn.PixelShuffle(upscale_factor=3)
# b = torch.rand((2,36,8,8))
# a(b).shape
# dir(a)

# a = GeneratorResNet()
# b = torch.rand((2, 3, 16, 16))
# a(b).shape

# Discriminator ResNet: Real? vs Fake? -----------------------------------------------------------------------------------
class Discriminator(nn.Module):     # real or fake // binary classification 마지막 activation으로 sigmoid를 사용해주어야함 (Vanilla GAN)
    def __init__(self, input_shape):    # loss 도 BCELoss를 사용해야함 → GAN Train시 불안정하여 MSE Loss를사용(LSGAN 구조)
        super(Discriminator, self).__init__()   # output size가 1×1이 아님, Patch로 되어 있어서 PatchGAN구조이다.

        self.input_shape = input_shape
        in_channels, in_height, in_width = self.input_shape
        patch_h, patch_w = int(in_height / 2 ** 4), int(in_width / 2 ** 4)  # patch_size:  h / 16, w / 16
        self.output_shape = (1, patch_h, patch_w)   # Patch 형태로 되어있음.

        def discriminator_block(in_filters, out_filters, first_block=False):
            layers = []
            layers.append(nn.Conv2d(in_filters, out_filters, kernel_size=3, stride=1, padding=1))
            if not first_block:
                layers.append(nn.BatchNorm2d(out_filters))
            layers.append(nn.LeakyReLU(0.2, inplace=True))  # ReLU vs. LeakyReLU  // LeakyReLU: 실험적으로 Discriminator에 Leaky ReLU를 써야 학습이 잘됨
            layers.append(nn.Conv2d(out_filters, out_filters, kernel_size=3, stride=2, padding=1))
            layers.append(nn.BatchNorm2d(out_filters))
            layers.append(nn.LeakyReLU(0.2, inplace=True))
            return layers

        layers = []
        in_filters = in_channels
        for i, out_filters in enumerate([64, 128, 256, 512]):
            layers.extend(discriminator_block(in_filters, out_filters, first_block=(i == 0)))
            in_filters = out_filters

        layers.append(nn.Conv2d(out_filters, 1, kernel_size=3, stride=1, padding=1))

        self.model = nn.Sequential(*layers)

    def forward(self, img):
        return self.model(img)

# b = torch.rand((2, 3, 64, 64))
# a = Discriminator(b.shape[1:])
# a(b).shape

os.makedirs("images", exist_ok=True)
os.makedirs("saved_models", exist_ok=True)

epoch = 0       # epoch to start training from
n_epochs = 10   # number of epochs of training

dataset_name = 'img_align_celeba' # name of dataset

batch_size = 4      # size of the batches
lr = 2e-4           # learning rate
b1 = 0.5            # decay of first order momentum of gradient
b2 = 0.999          # decay of first order momentum of gradient
decay_epoch = 100   # epoch from which to start lr decay

hr_height = 256
hr_width = 256
channels = 3
sample_interval = 100
checkpoint_interval = 10

cuda = torch.cuda.is_available()
if cuda:
    n_cpu = 8 # number of cpu threads to use during batch generation
else:
    n_cpu = 0

hr_shape = (hr_height, hr_width)    # (256, 256)

# Initialize generator and discriminator
generator = GeneratorResNet()
# generator
discriminator = Discriminator(input_shape=(channels, *hr_shape))
# discriminator
feature_extractor = FeatureExtractor()  # perceptual loss
# feature_extractor

# Set feature extractor to inference mode
feature_extractor.eval()    # dropout & batchnorm - training 단계에서만 On되어야함, Test일때 off

# Losses Function 정의
criterion_GAN = torch.nn.MSELoss()      # Loss Function
criterion_content = torch.nn.L1Loss()   # L1 Loss 큰 차이는 없고, 그래프가 V자형 직선 / L2 Loss 그래프가 U자형 곡선: 미분시 차이가 있음, U자형은 0에가까울수록 기울기가 0에 가까워짐

if cuda:
    generator = generator.cuda()
    discriminator = discriminator.cuda()
    feature_extractor = feature_extractor.cuda()
    criterion_GAN = criterion_GAN.cuda()
    criterion_content = criterion_content.cuda()

if epoch != 0:
    # Load pretrained models
    generator.load_state_dict(torch.load("saved_models/generator_%d.pth"))          # (nn.Module) Load_State_Dict(torch.load(file_path))
    # torch.load하면, dict 형태로 load해줌, 그리고 이 dict에 정의된 weight를 모델에 정의해주는게 load_state_dict
    # ['cov1': weight (torch.cuda.Tensor), ...]
    # Key와 이름이 같아야 weight를 적용.
    discriminator.load_state_dict(torch.load("saved_models/discriminator_%d.pth"))  #

# Optimizers
optimizer_G = torch.optim.Adam(generator.parameters(), lr=lr, betas=(b1, b2))       # Generator
optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=lr, betas=(b1, b2))   # Discriminator

Tensor = torch.cuda.FloatTensor if cuda else torch.Tensor   # cuda일때는 cuda tensor가 되도록

dataloader = DataLoader(
    ImageDataset("data/%s" % dataset_name, hr_shape=hr_shape),
    batch_size=batch_size,
    shuffle=True,
    num_workers=n_cpu,
)

# Sample Image
data_example = next(iter(dataloader))
len(data_example)
data_example.keys()
data_example['lr'].shape        # low resolution
data_example['hr'].shape        # high resolution

sample_index = 1
plt.figure(figsize=(10,20))
plt.subplot(121)
plt.imshow(np.transpose(data_example['lr'].numpy(), [0,2,3,1])[sample_index,:])
plt.subplot(122)
plt.imshow(np.transpose(data_example['hr'].numpy(), [0,2,3,1])[sample_index,:])
plt.show()

# Test Code
def tensor_to_image(tensor):
    return np.transpose(tensor.detach().to('cpu').numpy(), [0, 2, 3, 1])

a = data_example["lr"][[0],:]
b = data_example["hr"][[0],:]
g_a = generator(a.to(device))

plt.figure(figsize=(10,30))
plt.subplot(131)
plt.imshow(tensor_to_image(a).squeeze())
plt.subplot(132)
plt.imshow(tensor_to_image(b).squeeze())
plt.subplot(133)
plt.imshow(tensor_to_image(g_a).squeeze())
plt.show()

# d_a = discriminator(g_a)
# d_a

# ----------
#  Training
# ----------
# n_epochs = 1
for epoch in range(epoch, n_epochs):
    for i, imgs in enumerate(dataloader):

        # Configure model input
        # imgs_lr = imgs["lr"]    # shape: (4, 3, 64, 64)
        # imgs_hr = imgs["hr"]    # shape: (4, 3, 256, 256)
        imgs_lr = Variable(imgs["lr"].type(Tensor))     # 요즘 torch에서는 Variable로 안감싸도됩니다.
        imgs_hr = Variable(imgs["hr"].type(Tensor))     # require_grad를 위해서 Variable 해줬던 것...

        # Adversarial ground truths
        valid = Variable(Tensor(np.ones((imgs_lr.size(0), *discriminator.output_shape))), requires_grad=False)
        fake = Variable(Tensor(np.zeros((imgs_lr.size(0), *discriminator.output_shape))), requires_grad=False)

        # ------------------
        #  Train Generators
        # ------------------

        optimizer_G.zero_grad()

        # Generate a high resolution image from low resolution input
        gen_hr = generator(imgs_lr)

        # Adversarial loss
        loss_GAN = criterion_GAN(discriminator(gen_hr), valid)  # MSE Loss

        # Content loss
        gen_features = feature_extractor(gen_hr)
        real_features = feature_extractor(imgs_hr)
        loss_content = criterion_content(gen_features, real_features.detach())
        # detach(): tensor가 network돌면서 computational graph 만듦.
        #           뒤에 real_features에서는 gradient계산이 필요없기 때문에 detach() 해줌

        # Total loss
        loss_G = loss_content + 1e-3 * loss_GAN     # (hyper-parameter) 1:1e-3, (MSE Loss가 빠져있는데 이 Loss를 넣어도됨)

        loss_G.backward()
        optimizer_G.step()

        # ---------------------
        #  Train Discriminator
        # ---------------------

        optimizer_D.zero_grad()

        # Loss of real and fake images
        loss_real = criterion_GAN(discriminator(imgs_hr), valid)
        loss_fake = criterion_GAN(discriminator(gen_hr.detach()), fake)
        # generator 거친 computarional graph는 필요없음. 왜냐면 discriminator만 update할 거라서
        # back prop시 discriminator를 업데이트 할꺼면 discriminator까지만 가면됨. 더 뒤로갈 필요없음(더 뒤의 computaional그래프는 버리면됨)

        # (Forward) image → Generator → Discriminator → logit
        # (Backward) logit → Discriminator → 필요 ×

        # Total loss
        loss_D = (loss_real + loss_fake) / 2

        loss_D.backward()
        optimizer_D.step()

        # --------------
        #  Log Progress
        # --------------

        sys.stdout.write(
            "[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]\n"
            % (epoch, n_epochs, i, len(dataloader), loss_D.item(), loss_G.item())
        )

        batches_done = epoch * len(dataloader) + i
        if batches_done % sample_interval == 0:
            # Save image grid with upsampled inputs and SRGAN outputs
            imgs_lr = nn.functional.interpolate(imgs_lr, scale_factor=4)        # Resize해서 HR이미지와 같이 visualize 하는 부분
            gen_hr = make_grid(gen_hr, nrow=1, normalize=True)
            imgs_lr = make_grid(imgs_lr, nrow=1, normalize=True)
            img_grid = torch.cat((imgs_lr, gen_hr), -1)
            save_image(img_grid, "images/%d.png" % batches_done, normalize=False)

    if checkpoint_interval != -1 and epoch % checkpoint_interval == 0:
        # Save model checkpoints
        torch.save(generator.state_dict(), "saved_models/generator_%d.pth" % epoch)             # Generator 모델 저장
        torch.save(discriminator.state_dict(), "saved_models/discriminator_%d.pth" % epoch)     # Discriminator 모델 저장
        # (Load) torch.load → model.load_state_dict()
        # (Save) torch.state_dict() → torch.save
        # ※ model.load_state_dict(): weight를 dictionary형태로 뽑아주는 method

# load_generator = torch.load( "saved_models/generator_%d.pth" % epoch)
torch.save(generator, 'saved_models/SRGAN_generator.pt')
torch.save(discriminator, 'saved_models/SRGAN_discriminator.pt')

# SRCNN TEST ================================================================================================
weights_file = "checkpoint/srcnn_x3.pth"
image_file = "sample/zebra.bmp"
# image_file = "sample/ppt3.bmp"
scale = 3

cudnn.benchmark = True
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

model = SRCNN().to(device)

state_dict = model.state_dict()
for n, p in torch.load(weights_file, map_location=lambda storage, loc: storage).items():
    if n in state_dict.keys():
        state_dict[n].copy_(p)
    else:
        raise KeyError(n)

model.eval()

image = pil_image.open(image_file).convert('RGB')

image_width = (image.width // scale) * scale
image_height = (image.height // scale) * scale
image = image.resize((image_width, image_height), resample=pil_image.BICUBIC)
image = image.resize((image.width // scale, image.height // scale), resample=pil_image.BICUBIC)
image = image.resize((image.width * scale, image.height * scale), resample=pil_image.BICUBIC)
image.save(image_file.replace('.', '_bicubic_x{}.'.format(scale)))

image = np.array(image).astype(np.float32)
ycbcr = convert_rgb_to_ycbcr(image)

y = ycbcr[..., 0]
y /= 255.
y = torch.from_numpy(y).to(device)
y = y.unsqueeze(0).unsqueeze(0)

with torch.no_grad():
    preds = model(y).clamp(0.0, 1.0)

psnr = calc_psnr(y, preds)
print('PSNR: {:.2f}'.format(psnr))

preds = preds.mul(255.0).cpu().numpy().squeeze(0).squeeze(0)

output = np.array([preds, ycbcr[..., 1], ycbcr[..., 2]]).transpose([1, 2, 0])
output = np.clip(convert_ycbcr_to_rgb(output), 0.0, 255.0).astype(np.uint8)
output = pil_image.fromarray(output)
output.save(image_file.replace('.', '_srcnn_x{}.'.format(scale)))


import numpy as np
import pandas as pd
from pandas.tseries.offsets import DateOffset

def get_asset_data(underlyings, date_str):
    filt_unds = []
    spot_prices = []

    # TODO: потом сделать подсос с yahoo finance
    for und in underlyings:
        und_df = pd.read_csv(f'data/underlying/{und}.csv', sep=';',
                             parse_dates=['date'])

        und_df = \
            und_df.loc[(und_df['date'] <= date_str) &
                       (und_df['date'] >= pd.to_datetime(date_str)
                        - DateOffset(months=12))].set_index('date')

        filt_und = und_df[:date_str]
        spot_prices.append(float(und_df.iloc[-1]))
        filt_unds.append(filt_und)

    spot_prices = np.array(spot_prices)

    df = pd.concat(filt_unds, axis=1)

    # log-returns
    daily_returns = np.log(df) - np.log(df).shift(1)
    daily_returns.dropna(axis=0, inplace=True)
    # df.pct_change().dropna(axis=0)
    # monthly_returns = df.resample('M').first().pct_change().dropna(axis=0)

    ann_drift = (df.iloc[-1] / df.iloc[0]) ** (12/12) - 1
    # ann_vol = monthly_returns.std(axis=0) * 12 ** 0.5
    ann_vol = daily_returns.std(axis=0) * np.sqrt(daily_returns.shape[0])

    if len(underlyings) > 1:
        und_corr = daily_returns.corr()

        return spot_prices, ann_vol.to_numpy(), \
            und_corr.to_numpy(), ann_drift.to_numpy()

    return spot_prices[0], ann_vol[0], None, ann_drift[0]

def get_risk_free_rate_data(currency, date_str):
    df = pd.read_csv(f'data/currencies/{currency}.csv', sep=';',
                     parse_dates=['date'], dayfirst=True, index_col=0)

    row = df[:date_str].iloc[-1].dropna() / 100
    tenors = np.array([int(x) for x in row.index]) / 12

    return tenors, row.to_numpy()

def gbm_one_asset(S0, mu, sigma,
                  T, dt, num_simulations=10):

    N = int(T / dt)

    S = np.zeros((N + 1, num_simulations))
    S[0] = S0

    dW = np.random.normal(0.0, 1.0, (N, num_simulations))

    # cumprod?
    for i in range(N):
        S[i + 1] = S[i] * \
                   np.exp((mu - sigma ** 2 / 2) * dt + sigma * np.sqrt(dt) * dW[i])

    S = S[1:, :].T
    S = S = S[:, np.newaxis, :]

    return S

def gbm_multi_asset(S0, mu, sigma, rho,
                    T, dt, num_simulations = 10):

    N = int(T / dt)

    L = np.linalg.cholesky(rho) if rho is not None else np.eye(S0.shape[0])

    S = np.zeros((N+1, S0.shape[0], num_simulations))
    S[0] = S0.reshape(-1, 1)

    dW = np.random.normal(0.0, 1.0, (N, num_simulations, S0.shape[0]))
    dW_correlated = np.matmul(dW, L.T)

    # cumprod?
    for i in range(N):
        it = (mu - sigma**2 / 2) * dt + sigma * np.sqrt(dt) * dW_correlated[i]
        S[i+1] = S[i] * np.exp(it).T

    S = S[1:, :, :].T
    return S

def get_price_paths(spot_prices, ann_drift, ann_vol,
                     correlations, term, observation_period, num_simulations, method):

    if method == 'stochastic' and type(spot_prices) in (int, float):
        pass
        # generate by heston model

    elif method == 'constant' and type(spot_prices) == np.ndarray:
        price_paths = gbm_multi_asset(spot_prices, ann_drift, ann_vol, correlations,
                                      term, observation_period, num_simulations)

    else:
        # generate with one asset gbm
        price_paths = gbm_one_asset(spot_prices, ann_drift, ann_vol,
                                    term, observation_period, num_simulations)

    return price_paths


import backtrader as bt
import sqlite3
import pandas as pd
import datetime

# 투자전략 만들기
class TestStrategy(bt.Strategy):

    def log(self, txt, dt=None):
        '''전략의 로그 남기기'''
        dt = dt or self.datas[0].datetime.date(0)
        print('%s, %s' % (dt.isoformat(), txt))
        # if self.params.printlog or doprint:
        #     dt = dt or self.datas[0].datetime.date(0)
        #     print('%s, %s' % (dt.isoformat(), txt))

    def __init__(self):
        # 변수 모음 ==============================================================================
        self.dataclose = self.datas[0].close
        self.dataopen = self.datas[0].open

        self.order = None
        self.buyprice = None
        self.buycomm = None
        # ======================================================================================

        # 지표 추가 ========================================================================================
        self.sma = bt.indicators.MovingAverageSimple(self.datas[0], period=20)
        # =================================================================================================

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            return

        # 주문 완료됐는지 체크
        if order.status in [order.Completed]:
            if order.isbuy():
                self.log('매수 체결| 매수가: %.2f, 비용: %.2f, 수수료: %.2f' %
                         (order.executed.price, order.executed.value, order.executed.comm))
            elif order.issell():
                self.log('매도 체결| 매도가: %.2f, 비용: %.2f, 수수료: %.2f' %
                         (order.executed.price, order.executed.value, order.executed.comm))

            self.bar_executed = len(self)

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('주문 취소/거절')

        self.order = None

    def notify_trade(self, trade):
        if not trade.isclosed:
            return

        self.log('총손익: %.2f, 순손익: %.2f' %
                 (trade.pnl, trade.pnlcomm))

    def next(self):
        self.log('시가: %.2f 종가: %.2f' % (self.dataopen[0], self.dataclose[0]))

        # 진행 중인 주문이 있는지 확인
        if self.order:
            return

        # 포지션 보유 중인지 확인
        if not self.position:
            # 3일 연속 종가 하락하면 Buy 시그널 =======================================================================================
            ## 직전 일자 데이터 없으면 가장 최근데이터 가져옴 >> data length 확인 필요
            if len(self) > 3:
                if self.dataclose[0] < self.dataclose[-1]:
                    if self.dataclose[-1] < self.dataclose[-2]:
                        self.log('매수 시그널: %.2f (전일종가: %.2f | 2일전종가: %.2f)' % (self.dataclose[0], self.dataclose[-1], self.dataclose[-2]))
                        self.order = self.buy(coc=True)

        else:
            # 보유 후 N일 경과하면 매도
            if len(self) >= (self.bar_executed + 5):
                self.log('매도 시그널: %.2f' % self.dataclose[0])

                self.order = self.sell(coc=False)

if __name__ == '__main__':
    cerebro = bt.Cerebro()

    # 전략 추가 ===================================================================================
    cerebro.addstrategy(TestStrategy)

    # 데이터 가져오기 ===============================================================================
    from_date = "20200801"
    to_date = ''
    data_path = 'E:/DB/cmp_ohlcv.db'
    ## 시총 상위 200개 종목
    sql = "SELECT * FROM cmp_master WHERE Cmp_type=1 AND List_date < " + from_date + " ORDER BY Mkt_cap DESC LIMIT 200"
    con = sqlite3.connect(data_path)
    code_list = pd.read_sql(sql, con)
    code_list = code_list['Code']
    # ============================================================================================

    # 데이터 Cerebro에 입력 =========================================================================
    sql = "SELECT DAY, OPEN, HIGH, LOW, CLOSE, VOLUME  FROM T_STK_DATA WHERE CMP_CD = '{0}' AND DAY > " + from_date + " ORDER BY DAY ASC"

    code = 'A069500'
    data = pd.read_sql(sql.format(code), con, index_col='DAY', parse_dates=['DAY'])
    data = bt.feeds.PandasData(dataname=data, name=code)
    cerebro.adddata(data)
    # ============================================================================================

    # 초기현금, 수수료, Sizer 세팅 ===================================================================
    cerebro.broker.set_cash(10000000)
    cerebro.broker.setcommission(0.015)
    cerebro.addsizer(bt.sizers.FixedSize, stake=10)
    cerebro.broker.set_coc(True)
    # ============================================================================================

    print('기초 포트폴리오 가치: %.2f' % cerebro.broker.getvalue())

    cerebro.run()

    print('기말 포트폴리오 가치: %.2f' % cerebro.broker.getvalue())

    cerebro.plot()


import random
import numpy as np
import matplotlib.pyplot as plt
import math
from scipy import stats

plt.style.use('ggplot')

def direct_method(probs, samples=10000, plot=False):
    """
    Function for direct sampling method
    """

    sample_values = []

    for i in range(samples):
        sample = random.random()

        for j in range(len(probs)):
            if sample > sum(probs[0:j+1]):
                continue
            else:
                sample_values.append(j+1)
                break

    if plot:
        plt.hist(sample_values, bins=len(probs), rwidth=0.75)
        plt.title('Direct sample method')
        plt.show()

    return sample_values

def rejection_method(probs, sample_values=10000, plot=False):
    """
    Function for rejection sampling method
    """

    # Compute c
    k = len(probs)
    c = max(k * probs)

    num_reject = 0
    samples = []

    while len(samples) < sample_values:
        # get U1, U2
        rand1 = random.random()
        rand2 = random.random()

        # Get indicator variable
        indicator = math.floor(k * rand1) + 1

        # see if we should reject
        if rand2 <= probs[indicator-1]/c:
            samples.append(indicator)

        # Reject and try once more
        else:
            num_reject += 1
            continue

    if plot:
        plt.hist(samples, bins=k, rwidth=0.75)
        plt.title(
            f'Rejection method, {num_reject} rejected')
        plt.show()

    return samples

def alias_method(probs, num_samples=10000, plot=False):
    """
    Function for alias sampling method
    """

    # Generate F and L
    k = len(probs)

    L = list(range(1, k+1))
    F = k * probs

    # Generate g and s
    g = list(np.array(L)[(F >= 1)])
    s = list(np.array(L)[(F <= 1)])

    while len(s) > 0:

        i, j = g[0], s[0]

        L[j-1] = i
        F[i-1] = F[i-1] - (1-F[j-1])

        if F[i-1] < 1 - 1e-5:

            # Remove first element
            g.pop(0)

            # Add i to S
            s.append(i)

        s.pop(0)

    samples = []

    while len(samples) < num_samples:
        rand1 = random.random()
        rand2 = random.random()

        indicator = math.floor(k * rand1) + 1

        if rand2 <= F[indicator-1]:
            samples.append(indicator)
        else:
            samples.append(L[indicator-1])

    if plot:
        plt.hist(samples, bins=len(probs), rwidth=0.75)
        plt.title('Alias Method')
    plt.show()

    return samples

# Probabilities that will be used for simulation
probs = np.array([7/48, 5/48, 1/8, 1/16, 1/4, 5/16])

plt.bar(list(range(1, len(probs) + 1)), probs)
plt.title('Probabilites')
plt.show()

num_samples = 10000

# Direct method
samples_direct = direct_method(probs, samples=num_samples, plot=True)
frequencies_direct = np.bincount(samples_direct)[1:]

chisq_direct, pval_direct = stats.chisquare(
    frequencies_direct, probs*num_samples)
print(f'Direct Test statistic {chisq_direct}')
print(f'Direct pval {pval_direct}')

# Rejection method
samples_reject = rejection_method(probs, sample_values=num_samples, plot=True)
frequencies_reject = np.bincount(samples_reject)[1:]

chisq_reject, pval_reject = stats.chisquare(
    frequencies_reject, probs*num_samples)
print(f'Reject Test statistic {chisq_reject}')
print(f'Reject pval {pval_reject}')

# Alias method
samples_alias = alias_method(probs, num_samples=num_samples, plot=True)
frequencies_alias = np.bincount(samples_alias)[1:]

chisq_alias, pval_alias = stats.chisquare(frequencies_alias, probs*num_samples)
print(f'Alias Test statistic {chisq_alias}')
print(f'Alias pval {pval_alias}')


import logging as lg
import os
import re
from aztest.common import to_list
from aztest.errors import AzTestError
logger = lg.getLogger(__name__)

DEFAULT_WHITELIST_FILE = "lmbr_test_whitelist.txt"
DEFAULT_BLACKLIST_FILE = "lmbr_test_blacklist.txt"

def get_default_whitelist():
    """Returns the default whitelist file path if it exists, otherwise empty string

    :return: whitelist path or empty string
    :rtype: str
    """
    if os.path.exists(os.path.abspath(DEFAULT_WHITELIST_FILE)):
        return os.path.abspath(DEFAULT_WHITELIST_FILE)
    else:
        return ""

def get_default_blacklist():
    """Returns the default blacklist file path if it exists, otherwise empty string

    :return: blacklist path or empty string
    :rtype: str
    """
    if os.path.exists(os.path.abspath(DEFAULT_BLACKLIST_FILE)):
        return os.path.abspath(DEFAULT_BLACKLIST_FILE)
    else:
        return ""

class FileApprover(object):
    """Class for compiling and validating the set of files that are allowed to be tested"""
    whitelist = None
    blacklist = None

    def __init__(self, whitelist_files=None, blacklist_files=None):
        self.make_whitelist(whitelist_files)
        self.make_blacklist(blacklist_files)

    def make_whitelist(self, whitelist_files):
        """Make the whitelist from a file or list of files

        It is assumed that if no whitelist is provided, then we will scan everything not in a blacklist.

        :param whitelist_files: path to a whitelist file or list of paths
        """
        if whitelist_files:
            for whitelist_file in to_list(whitelist_files):
                self.add_whitelist(whitelist_file)
        else:
            self.whitelist = None

    def make_blacklist(self, blacklist_files):
        """Make the blacklist from a file or list of files

        It is assumed that is no blacklist is provided, then we will scan everything in a directory or what is
        whitelisted.

        :param blacklist_files: path to a blacklist file or list of paths
        """
        if blacklist_files:
            for blacklist_file in to_list(blacklist_files):
                self.add_blacklist(blacklist_file)
        else:
            self.blacklist = None

    def add_whitelist(self, whitelist_file):
        """Add the file of patterns to the whitelist

        :param whitelist_file: path to a whitelist file
        """
        if whitelist_file:
            patterns = self.get_patterns_from_file(whitelist_file)
            if patterns:
                self.whitelist = patterns | (self.whitelist or set())

    def add_blacklist(self, blacklist_file):
        """Add the file of patterns to the blacklist

        :param blacklist_file: path to a blacklist file
        """
        if blacklist_file:
            patterns = self.get_patterns_from_file(blacklist_file)
            if patterns:
                self.blacklist = patterns | (self.blacklist or set())

    def add_whitelist_patterns(self, patterns):
        """Add patterns to the whitelist

        :param patterns: regular expression pattern or list of patterns to add to whitelist
        """
        self.whitelist = set(to_list(patterns)) | (self.whitelist or set())

    def add_blacklist_patterns(self, patterns):
        """Add patterns to the blacklist

        :param patterns: regular expression pattern or list of patterns to add to blacklist
        """
        self.blacklist = set(to_list(patterns)) | (self.blacklist or set())

    def is_approved(self, file_name):
        """Checks to see if file_name is in the whitelist and not the blacklist

        :param file_name: name or path of the file to check
        :return: true if file_name in whitelist or whitelist is None and not in blacklist or blacklist is None,
        else false
        """
        return self.is_whitelisted(file_name) and not self.is_blacklisted(file_name)

    def is_whitelisted(self, file_name):
        """Checks to see if file_name is in the whitelist

        :param file_name: name or path of the file to check
        :return: true if file_name in whitelist or whitelist is None, else false
        """
        return True if not self.whitelist else self.is_in_list(file_name, self.whitelist, 'whitelist')

    def is_blacklisted(self, file_name):
        """Checks to see if file_name is in the blacklist

        :param file_name: name or path of the file to check
        :return: true if file_name not in blacklist or blacklist is None, else false
        """
        return False if not self.blacklist else self.is_in_list(file_name, self.blacklist, 'blacklist')

    @staticmethod
    def get_patterns_from_file(pattern_file):
        """Returns set of patterns from pattern_file if pattern_file is valid, otherwise returns None

        :param pattern_file: path of the whitelist or blacklist file to get patterns from
        :return: the set of patterns from the file, or None if pattern_file is invalid
        :rtype: set
        """
        if not pattern_file or not os.path.exists(pattern_file):
            raise AzTestError("Invalid module/pattern file path: {}".format(pattern_file))
        logger.info("Using module/pattern file: {}".format(pattern_file))
        with open(pattern_file, 'r') as f:
            patterns = set((line.strip() for line in f))
            return patterns

    @staticmethod
    def is_in_list(file_name, patterns, list_name='CMD'):
        """Checks file_name against the set of patterns using regex

        For the curious, we do not compile the patterns ahead of time because the 're' module uses an internal cache
        that stores compiled patterns, and it is large enough (100) that we shouldn't need to worry about patterns being
        recompiled during a scan.

        :param file_name: name or path of the file to check
        :param patterns: the set of patterns to use as regular expressions
        :param list_name: the name of the list being checked (if applicable, defaults to CMD)
        :return: true if there is a match, otherwise false
        """
        if not file_name or not patterns:
            return False
        logger.debug("Checking file: {}".format(file_name))
        for pattern in patterns:
            full_pattern = r"^.*[/\\]{}(\.dll|\.exe|\.dylib)?$".format(pattern)
            match = re.match(full_pattern, file_name, re.IGNORECASE)
            if match:
                logger.debug("Found file {} in {} patterns.".format(file_name, list_name))
                return True
        return False


import logging
import struct
from rdflib import BNode, Literal, URIRef
from aiohttp_rdf4j.utils import Statement

logger = logging.getLogger(__name__)

# Constants
MAGIC_NUMBER = b'BRDF'

FORMAT_VERSION = 1

#record types

NAMESPACE_DECL = 0
STATEMENT = 1
COMMENT = 2
VALUE_DECL = 3
ERROR = 126
END_OF_DATA = 127

#value types

NULL_VALUE = 0
URI_VALUE = 1
BNODE_VALUE = 2
PLAIN_LITERAL_VALUE = 3
LANG_LITERAL_VALUE = 4
DATATYPE_LITERAL_VALUE = 5
VALUE_REF = 6

class BinaryRDFParser:

    def __init__(self, stream):
        self.declaredValues = dict()
        self.namespaces = dict()
        #Utf8Decoder = codecs.getincrementaldecoder('utf-16be')
        #self.decoder = Utf8Decoder()
        self.stream = stream
        self.buffer = b""
        #self.vars = None #Dummy attribute so that query can ask for vars in construct
        #self.cache = collections.deque()

    async def read_(self,n, t=None):
        while len(self.buffer)<n:
            chunk, eos = await self.stream.readchunk()
            self.buffer+=chunk
        res, self.buffer = self.buffer[:n], self.buffer[n:]
        return struct.unpack(t, res) if t else res

    async def parse(self):
        mn= await self.read_(4)
        if mn != MAGIC_NUMBER:
            ##print ("No Magic")
            raise IOError("Bad Magic Number")
            #return

        fv, = await self.read_(4,"!i")
        #print("fv", fv)
        if  fv != FORMAT_VERSION:
            raise IOError("Wrong Version")
            #return

        while True:
            recordType, = await self.read_(1,"!b")
            #print("in while", recordType)
            if recordType == END_OF_DATA:
                break
            elif recordType == STATEMENT:
                yield await self.readStatement()

            elif recordType == VALUE_DECL:
                await self.readValueDecl()
            elif recordType == NAMESPACE_DECL:
                #print("ns")
                await self.readNamespaceDecl()
            elif recordType == COMMENT:
                await self.readComment()
            else:
                raise IOError("Parse Error for recordType %s", recordType)
                break

    async def readNamespaceDecl(self):
        #print("in namespace")
        prefix = await self.readString()
        #print("prefix", prefix)
        namespace = await self.readString()
        #print("namespace", namespace)
        self.namespaces[prefix] = namespace

    async def readComment(self):
        comment = await self.readString()

    async def readValueDecl(self):
        #self.stream_caching(4)
        id, = await self.read_(4, "!i")
        v = await self.readValue()
        self.declaredValues[id] = v

    async def readStatement(self):

        subj = await self.readValue()
        #print("subj", subj)
        pred = await self.readValue()
        #print("pred", pred)
        obj = await self.readValue()
        #print("obj", obj)
        ctx = await self.readValue()
        #print("ctx", ctx)

        #print ("ii", subj,pred,obj, ctx)
        return Statement(subj,pred,obj,ctx)

        #return ((subj, pred, obj), ctx)

    async def readValue(self):

        valueType, = await self.read_(1, "!b")
        ##print(valueType)
        if valueType == NULL_VALUE:
            return None
        elif valueType == VALUE_REF:
            return await self.readValueRef()
        elif valueType == URI_VALUE:
            return await self.readUri()
        elif valueType == BNODE_VALUE:
            return await self.readBNode()
        elif valueType == PLAIN_LITERAL_VALUE:
            return await self.readPlainLiteral()
        elif valueType == LANG_LITERAL_VALUE:
            return await self.readLangLiteral()
        elif valueType == DATATYPE_LITERAL_VALUE:
            return await self.readDatatypeLiteral()
        else:
            return None

    async def readValueRef(self):
        id, = await self.read_(4, "!i")
        return self.declaredValues[id]

    async def readUri(self):
        return URIRef(await self.readString())

    async def readBNode(self):
        return BNode(await self.readString())

    async def readPlainLiteral(self):
        return Literal(await self.readString())

    async def readLangLiteral(self):
        label = await self.readString()
        language = await self.readString()
        return Literal(label, lang=language)

    async def readDatatypeLiteral(self):
        label = await self.readString()
        datatype = await self.readString()
        dt = URIRef(datatype)
        return Literal(label, datatype=dt)

    async def readString(self):
        ##print(x)
        length, = await self.read_(4, "!i")
        stringBytes = length << 1
        ##print(self.stream.getbuffer().nbytes, stringBytes, self.stream.tell())
        #self.stream_caching(length)
        _t = await self.read_(stringBytes)
        string = _t.decode("utf-16be")
        #string = self.decoder.decode(_t)
        ##print (string)
        return string

    def __aiter__(self):
        return self.parse()


#Gerador de nomes

#'bcdfghjklmnpqrstvwxyz'
#'aeiou'

import random
class NameGenerator:
    def __init__(self, qtdNomesMax = 5, vogs='aiiiiooou', cons='bcdfghjklmnpqrstvwxyz', nameMinSize = 4, nameMaxSize = 7):
        self.qtdNome = qtdNomesMax
        self.vog = vogs
        self.con = cons
        self.qtdNomes = random.randint(nameMinSize, nameMaxSize)

    def genetateName(self):
        name = []
        each = bool(False)

        for i in range(self.qtdNomes):
            nameSize = random.randint(3, 7)
            for j in range(nameSize):

                if each:
                    vogSel = random.randint(0, len(self.vog) - 1)
                    name.append(self.vog[vogSel])
                else:
                    conSel = random.randint(0, len(self.con) - 1)
                    letterCons = self.con[conSel]

                    if letterCons == 'r' or letterCons == 's':
                        if random.randint(0, 2) == 2:
                            name.append(self.con[conSel])
                            name.append(self.con[conSel])
                        else:
                            name.append(self.con['*'])
                    else:
                        name.append(self.con[conSel])

                each = not each
            name.append(' ')

        name = (''.join(name)).title()
        return name

genNam = NameGenerator(5, 'bbbbcdfggghjklmmmmnnnnpqqrrrrrrrrsssssssttttvwxyz', 'aaaaaaaaaaaaaaaeeeeeeeeeeeeeiiiiiiooooouuu', 2, 7)
print(genNam.genetateName())

Introdução à programação com python/Testes/NameGenerator.py


from common import *

#by listing all the section, then traverse that list to compute the correct result
#to compare with the result of the algorithm

def brute_force(num):
    a = load_testcase(num)
    s=a[0]
    for i in range(0,len(a)) :
        for j in range(i,len(a)):
            b = a[i:j+1]
            if (s>sum(b)) :
                s=sum(b)
    return s
#write to the output
def brute_force_write_output():
    clear_output(0)
    for x in range(1,11):
        write_output(brute_force(x),0)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Dec  5 16:05:02 2019

@author: diaa
"""

z_dim = 2
X_dim = 784
N = 1000

from torch import nn
import torch.nn.functional as F
class Encoder(nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()
        self.lin1 = nn.Linear(X_dim, N)
        self.lin2 = nn.Linear(N, N)
        # Gaussian code (z)
        self.lin3gauss = nn.Linear(N, z_dim)

    def forward(self, x):
        x = F.dropout(self.lin1(x), p=0.2, training=self.training)
        x = F.relu(x)
        x = F.dropout(self.lin2(x), p=0.2, training=self.training)
        x = F.relu(x)
        xgauss = self.lin3gauss(x)

        return xgauss

class Decoder(nn.Module):
    def __init__(self):
        super(Decoder, self).__init__()
        self.lin1 = nn.Linear(z_dim, N)
        self.lin2 = nn.Linear(N, N)
        self.lin3 = nn.Linear(N, X_dim)

    def forward(self, x):
        x = self.lin1(x)
        x = F.dropout(x, p=0.2, training=self.training)
        x = F.relu(x)
        x = self.lin2(x)
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.lin3(x)
        return F.sigmoid(x)


# -*- coding:utf-8 -*-
import json
import csv
import pandas as pd
import config
from requests_oauthlib import OAuth1Session

CK = config.CONSUMER_KEY
CS = config.CONSUMER_SECRET
AT = config.ACCESS_TOKEN
ATS = config.ACCESS_TOKEN_SECRET
twitter = OAuth1Session(CK, CS, AT, ATS)

# search
def search(url, keyword, count):
    params = {"q": keyword, "count": count}
    response = twitter.get(url, params=params)
    if response.status_code == 200:
        search_result = json.loads(response.text)
        fields = ["created_at", "id", "text"]
        df = pd.DataFrame(search_result["statuses"], columns=fields)
        df.to_csv("result.csv", index=False, quoting=csv.QUOTE_ALL, encoding="utf-8-sig")
    else:
        print("ERROR: %d" % response.status_code)

if __name__ == '__main__':
    search(config.URL_SEARCH, config.SEARCH_KEYWORD, 100)


from django.http import HttpResponse
from django.template import RequestContext, loader

from polls.models import Products, TDD_LTE_Phases, WMP_Phases, DCM_Phases

def index_validator(table):
    try:
        name = table.objects.order_by('-product_id')[0]
    except IndexError:
        name = 'null'
    return name

def index(request):
    #Get latest data from each Product/Subcomponent
    latest_tdlte = TDD_LTE_Phases.objects.order_by('-product_id')[0]
    latest_wmp = WMP_Phases.objects.order_by('-product_id')[0]
    latest_dcm = index_validator(DCM_Phases)
    template = loader.get_template('polls/index.html')
    #Add latest data into context, and pass it to overview page.
    context = RequestContext(request, {
        'latest_tdlte': latest_tdlte,
        'latest_wmp':latest_wmp,
        'latest_dcm':latest_dcm,
    })
    return HttpResponse(template.render(context))

def detail(request, product_id):
    return HttpResponse("You're looking at Products %s." % product_id)

def results(request, products_id):
    response = "You're looking at the results of Products %s."
    return HttpResponse(response % products_id)

def vote(request, products_id):
    return HttpResponse("You're voting on Products %s." % products_id)

# Chart generation
from chartit import DataPool, Chart

def chart_view_with_parameter(request, chart_name):

    template = loader.get_template('polls/chart.html')

    #Step 1: Create a DataPool with the data we want to retrieve.
    tdlte = \
        DataPool(
           series=
            [{'options': {
               'source': eval(chart_name).objects.all()},
              'terms': [
                'id',
                'entity_build',
                'majorRegression',
                'neve']}
             ])

    #Step 2: Create the Chart object
    cht = Chart(
            datasource = tdlte,
            series_options =
              [{'options':{
                  'type': 'line',
                  'stacking': False},
                'terms':{
                  'id': [
                    'entity_build',
                    'majorRegression',
                    'neve']
                  }}],
            chart_options =
              {'title': {
                   'text': 'Entity build history for ' + chart_name},
               'xAxis': {
                    'title': {
                       'text': 'Product ID'}}})

    #Step 3: Send the chart object to the template.
    context = RequestContext(request, {
        'charts': cht,
    })

    return HttpResponse(template.render(context))

def chart_view(request):

    template = loader.get_template('polls/chart_selector.html')

    #Step 1: Create a DataPool with the data we want to retrieve.
    tdlte = \
        DataPool(
           series=
            [{'options': {
               'source': WMP_Phases.objects.all()},
              'terms': [
                'id',
                'entity_build',
                'majorRegression',
                'neve']}
             ])

    #Step 2: Create the Chart object
    cht = Chart(
            datasource = tdlte,
            series_options =
              [{'options':{
                  'type': 'line',
                  'stacking': False},
                'terms':{
                  'id': [
                    'entity_build',
                    'majorRegression',
                    'neve']
                  }}],
            chart_options =
              {'title': {
                   'text': 'Entity build history for WMP_Phases'},
               'xAxis': {
                    'title': {
                       'text': 'Product ID'}}})

    #Step 3: Send the chart object to the template.
    context = RequestContext(request, {
        'charts': cht,
    })

    return HttpResponse(template.render(context))


import pandas
from sklearn import linear_model
#reading the file
df = pandas.read_csv("cars.csv")

#make list of independent values, these are the x values
#the dependent values we will store these in the y value
X = df[['Weight', 'Volume']]
y  = df['CO2']

regr = linear_model.LinearRegression()
regr.fit(X,y)
#Now we have a regression object that are ready to predict CO2 values based on a car's weight and volume:
#predict the CO2 emission of a car where the weight is 2300kg, and the volume is 1300cm3:
predictedCo2 = regr.predict([[2300,1300]])
print(predictedCo2)

#this pritns the coeficient of the regression object
print(regr.coef_)

#We have predicted that a car with 1.3 liter engine, and a weight of 2300 kg,
#will release approximately 107 grams of CO2 for every kilometer it drives.


from flask import Flask, render_template, request, redirect
from database.get_connection import get_connection
from database.init_db import init_db
from database.insert_db import insert_db
import click
from flask.cli import with_appcontext
import database.user_controller as user_controller
import database.providers_controller as providers_controller
import database.movies_controller as movies_controller
import database.rentals_controller as rentals_controller
import database.ratings_controller as ratings_controller
from datetime import date

app = Flask(__name__)
app.config.from_pyfile("config/config.py")

@app.route('/usuarios')
def get_users():
    users = user_controller.get_users()
    return render_template('users.html', users=users)

@app.route('/proveedores')
def get_providers():
    providers = providers_controller.get_providers()
    return render_template('providers.html', providers=providers)

@app.route('/peliculas')
def get_movies():
    movies = movies_controller.get_movies()
    return render_template('movies.html', movies=movies)

@app.route('/alquileres')
def get_rentals():
    rentals = rentals_controller.get_rentals()
    return render_template('rentals.html', rentals=rentals)

@app.route('/valoraciones')
def get_ratings():
    ratings = ratings_controller.get_ratings()
    return render_template('ratings.html', ratings=ratings)

@app.route('/methods/crear_pelicula', methods=['POST'])
def create_movie():
    title = request.form['title']
    category = request.form['category']
    description = request.form['description']
    provider = request.form['provider']
    units = request.form['units']

    # First primary key check
    exists = movies_controller.get_movie_by_title(title)
    if exists is not None:
        return render_template('add_movie.html',
                               error='Violación de integridad de clave primaria: La película ya ha sido agregada 🔑')

    # Second foreign key check
    exists = providers_controller.get_provider_by_company(provider)
    if exists is None:
        return render_template('add_movie.html',
                               error='Violación de integridad de clave foránea: La empresa proveedora no existe 🔐')

    movies_controller.insert_movie(title, category, description, provider, units)
    return redirect('/peliculas')

@app.route('/methods/crear_alquiler', methods=['POST'])
def create_rental():
    title = request.form['title']
    owner = request.form['owner']
    start = request.form['start']
    end = request.form['end']
    type_query = request.form['type']

    # First primary key check
    exists = rentals_controller.get_rental_by_title_and_owner(title, owner)
    if exists is not None:
        return render_template('add_rental.html',
                               error='Violación de integridad de clave primaria: El alquiler ya ha sido creado 🔑')

    # Second foreign key check (with movie)
    exists = movies_controller.get_movie_by_title(title)
    if exists is None:
        return render_template('add_rental.html',
                               error='Violación de integridad de clave foránea: La película introducida no existe 🔐')

    # Third foreign key check (with user)
    exists = user_controller.get_user_by_nick(owner)
    if exists is None:
        return render_template('add_rental.html',
                               error='Violación de integridad de clave foránea: El usuario introducido no existe 🔐')

    # We must check if the movie is available with the total units
    units = movies_controller.get_movie_units(title)

    if units <= 0:
        return render_template('add_rental.html',
                               error='Restricción del dominio del problema: No hay unidades disponibles 🧿')

    # We must check if introduced date is greater than the current date
    if start < date.today().strftime("%Y-%m-%d"):
        return render_template('add_rental.html',
                               error='Restricción del dominio del problema: La fecha de alquiler no puede ser menor que la fecha actual 🧿')

    # We must check that the return date is greater than the rental date
    if start > end:
        rental = rentals_controller.get_rental_by_title_and_owner(title, owner)
        return render_template('add_rental.html',
                               error='Restricción del dominio del problema: La fecha de devolución no puede ser menor que la fecha de alquiler 🧿')

    rentals_controller.insert_rental(title, owner, start, end, type_query)

    # We must decrease the final units of the movie
    movies_controller.decrease_units(title)

    return redirect('/alquileres')

@app.route('/methods/crear_valoracion', methods=['POST'])
def create_rating():
    title = request.form['title']
    author = request.form['author']
    score = request.form['rating']
    comment = request.form['comment']

    # First primary key check
    exists = ratings_controller.get_rating_by_title_and_author(title, author)
    if exists is not None:
        return render_template('add_rating.html',
                               error='Violación de integridad de clave primaria: La valoración ya ha sido creada 🔑')

    # Second foreign key check (with movie)
    exists = movies_controller.get_movie_by_title(title)
    if exists is None:
        return render_template('add_rating.html',
                               error='Violación de integridad de clave foránea: La película introducida no existe 🔐')

    # Third foreign key check (with user)
    exists = user_controller.get_user_by_nick(author)
    if exists is None:
        return render_template('add_rating.html',
                               error='Violación de integridad de clave foránea: El usuario introducido no existe 🔐')

    # We must check if the score is between 0 and 10
    if float(score) < 0 or float(score) > 10:
        return render_template('add_rating.html',
                               error='Restricción del dominio del problema: La puntuación debe estar entre 0 y 10 🧿')

    ratings_controller.insert_rating(title, score, comment, author)

    return redirect('/valoraciones')

@app.route('/methods/crear_proveedor', methods=['POST'])
def create_provider():
    company = request.form['company']
    nif = request.form['nif']
    phone = request.form['phone']

    exists = providers_controller.get_provider_by_company(company)
    if exists is not None:
        return render_template('add_provider.html',
                               error='Violación de integridad de clave primaria: La empresa ya ha sido creada 🔑')

    providers_controller.insert_provider(company, nif, phone)
    return redirect('/proveedores')

@app.route('/methods/eliminar_proveedor', methods=['POST'])
def delete_provider():
    company = request.form['company']
    providers_controller.delete_provider(company)
    return redirect('/proveedores')

@app.route('/methods/eliminar_pelicula', methods=['POST'])
def delete_movie():
    title = request.form['title']
    movies_controller.delete_movie(title)
    return redirect('/peliculas')

@app.route('/alquileres/agregar_alquiler')
def add_rental():
    return render_template('add_rental.html')

@app.route('/valoraciones/agregar_valoracion')
def add_rating():
    return render_template('add_rating.html')

@app.route('/proveedores/agregar_proveedor')
def add_provider():
    return render_template('add_provider.html')

@app.route('/usuarios/agregar_usuario')
def add_user():
    return render_template('add_user.html')

@app.route('/peliculas/agregar_pelicula')
def add_movie():
    return render_template('add_movie.html')

@app.route('/usuarios/editar_usuario/<string:id>')
def edit_user(id):
    user = user_controller.get_user_by_nick(id)
    return render_template('edit_user.html', user=user)

@app.route('/alquileres/editar_alquiler/<string:owner>/<string:title>')
def edit_rental(owner, title):
    rental = rentals_controller.get_rental_by_title_and_owner(title, owner)
    return render_template('edit_rental.html', rental=rental)

@app.route('/valoraciones/editar_valoracion/<string:author>/<string:title>')
def edit_rating(author, title):
    rating = ratings_controller.get_rating_by_title_and_author(title, author)
    return render_template('edit_rating.html', rating=rating)

@app.route('/peliculas/editar_pelicula/<string:id>')
def edit_movie(id):
    movie = movies_controller.get_movie_by_title(id)
    return render_template('edit_movie.html', movie=movie)

@app.route('/methods/editar_usuario', methods=['POST'])
def update_user():
    nick = request.form['nick']
    name = request.form['name']
    email = request.form['email']
    user_controller.update_user(nick, name, email)
    return redirect('/usuarios')

@app.route('/methods/editar_pelicula', methods=['POST'])
def update_movie():
    title = request.form['title']
    category = request.form['category']
    description = request.form['description']
    provider = request.form['provider']
    units = request.form['units']

    # Second foreign key check
    exists = providers_controller.get_provider_by_company(provider)
    if exists is None:
        movie = movies_controller.get_movie_by_title(title)

        return render_template('edit_movie.html',
                               error='Violación de integridad de clave foránea: La empresa proveedora no existe 🔐',
                               movie=movie)

    movies_controller.update_movie(title, category, description, provider, units)
    return redirect('/peliculas')

@app.route('/methods/editar_alquiler', methods=['POST'])
def update_rental():
    owner = request.form['owner']
    title = request.form['title']
    rental_date = request.form['start']
    return_date = request.form['end']
    type_query = request.form['type']

    # We must check if introduced date is greater than the current date
    if rental_date < date.today().strftime("%Y-%m-%d"):
        rental = rentals_controller.get_rental_by_title_and_owner(title, owner)
        return render_template('edit_rental.html',
                               error='Restricción del dominio del problema: La fecha de alquiler no puede ser menor que la fecha actual 🧿',
                               rental=rental)

    # We must check that the return date is greater than the rental date
    if rental_date > return_date:
        rental = rentals_controller.get_rental_by_title_and_owner(title, owner)
        return render_template('edit_rental.html',
                               error='Restricción del dominio del problema: La fecha de devolución no puede ser menor que la fecha de alquiler 🧿',
                               rental=rental)

    rentals_controller.update_rental(title, owner, rental_date, return_date, type_query)

    return redirect('/alquileres')

@app.route('/methods/editar_valoracion', methods=['POST'])
def update_rating():
    author = request.form['author']
    title = request.form['title']
    rating = request.form['rating']
    comment = request.form['comment']

    # We must check that the rating is between 0 and 10
    if float(rating) < 0 or float(rating) > 10:
        rating = ratings_controller.get_rating_by_title_and_author(title, author)

        return render_template('edit_rating.html',
                               error='Restricción del dominio del problema: La valoración debe estar entre 0 y 10 estrellas (⭐) 🧿',
                               rating=rating)

    ratings_controller.update_rating(title, rating, comment, author)

    return redirect('/valoraciones')

@app.route('/methods/crear_usuario', methods=['POST'])
def create_user():
    nick = request.form['nick']

    exists = user_controller.get_user_by_nick(nick)
    if exists is not None:
        return render_template('add_user.html', error='Violación de integridad de clave primaria: El nick ya existe 🔑')

    name = request.form['name']
    email = request.form['email']

    user_controller.insert_user(nick, name, email)
    return redirect('/usuarios')

@app.route('/methods/eliminar_usuario', methods=['POST'])
def delete_user():
    nick = request.form['nick']
    user_controller.delete_user(nick)
    return redirect('/usuarios')

@app.route('/methods/eliminar_alquiler', methods=['POST'])
def delete_rental():
    owner = request.form['owner']
    title = request.form['title']
    rentals_controller.delete_rental(title, owner)

    # We must increase the final units of the movie
    movies_controller.increase_units(title)

    return redirect('/alquileres')

@app.route('/methods/eliminar_valoracion', methods=['POST'])
def delete_rating():
    author = request.form['author']
    title = request.form['title']
    ratings_controller.delete_rating(title, author)
    return redirect('/valoraciones')

@app.route('/proveedores/editar_proveedor/<string:id>')
def edit_provider(id):
    provider = providers_controller.get_provider_by_company(id)
    return render_template('edit_provider.html', provider=provider)

@app.route('/methods/editar_proveedor', methods=['POST'])
def update_provider():
    company = request.form['company']
    nif = request.form['nif']
    phone = request.form['phone']
    providers_controller.update_provider(company, nif, phone)
    return redirect('/proveedores')

@click.command("init-db")
@with_appcontext
def init_db_command():
    """Clear the existing data and create new tables"""
    init_db()
    click.echo("Database initialized correctly! 🚀")
    insert_db()
    click.echo("All records in the database have been inserted correctly! 🔋")

app.cli.add_command(init_db_command)

@app.route('/restoreDatabaseEmpty')
def create_database():
    init_db()
    return str("The database has been successfully restored!")

@app.route('/restoreDatabaseInitialized')
def insert_records():
    # We must restore the database to its original state for insert the vanilla records
    init_db()
    insert_db()
    return str("The database has been successfully restored and initialized!")

@app.route('/')
def get_index():  # put application's code here
    return render_template('base.html')

if __name__ == '__main__':
    app.run(debug=True)


# python functionality for use with various python servers

import socket
import time

QUEUE_TIMEOUT_SECS = 4

HTTP_STATUS_OK = '200 OK'
HTTP_STATUS_TIMEOUT = '408 TIMEOUT'
HTTP_STATUS_BAD_REQUEST = '400 BAD REQUEST'

def simulated_file_read(elapsed_time_ms):
    time.sleep(elapsed_time_ms / 1000.0)  # seconds

def handle_request(request_text, server_name, receipt_timestamp):

    rc = HTTP_STATUS_BAD_REQUEST
    tot_request_time_ms = 0
    file_path = ''

    if len(request_text) > 0:
        start_processing_timestamp = time.time()
        queue_time_ms = start_processing_timestamp - receipt_timestamp
        queue_time_secs = queue_time_ms / 1000

        # has this request already timed out?
        if queue_time_secs >= QUEUE_TIMEOUT_SECS:
            print("timeout (queue)")
            rc = HTTP_STATUS_TIMEOUT
        else:
            line_tokens = request_text.split(' ')
            if len(line_tokens) > 1:
                request_method = line_tokens[0]
                args = line_tokens[1]
                fields = args.split(',')
                if len(fields) == 3:
                    #first_field = fields[0].strip('/')
                    #rc = int(first_field)
                    disk_read_time_ms = int(fields[1])
                    file_path = fields[2]
                    simulated_file_read(disk_read_time_ms)
                    rc = HTTP_STATUS_OK

        # total request time is sum of time spent in queue and the
        # simulated disk read time
        tot_request_time_ms = queue_time_ms + disk_read_time_ms

    # construct response and send back to client
    response_body = "%d,%s" % \
        (tot_request_time_ms, file_path)

    response_headers = 'HTTP/1.1 %s\n' % rc
    response_headers += "Server: %s\n" % server_name
    response_headers += "Content-Length: %d\n" % len(response_body)
    response_headers += "Connection: close\n"
    response_headers += "\n"

    return (response_headers, response_body)

def handle_socket_request(sock, server_name, receipt_timestamp):
    reader = sock.makefile('r')
    writer = sock.makefile('w')
    request_text = reader.readline()
    reader.close()

    headers, body = handle_request(request_text,
                                   server_name,
                                   receipt_timestamp)

    writer.write(headers)
    writer.write(body)
    writer.flush()

    writer.close()
    sock.close()


# -*- coding: utf-8 -*-
from __future__ import unicode_literals

from django.db import models, migrations

class Migration(migrations.Migration):

    dependencies = [
        ('api', '0004_housecode'),
    ]

    operations = [
        migrations.AddField(
            model_name='housecode',
            name='temperature_opentrv',
            field=models.CharField(default=None, max_length=6),
        ),
    ]


from collections import defaultdict, deque

class Graph:
    def __init__(self):
        self.graph = defaultdict(list)

    def add_edge(self, u, v):
        self.graph[u].append(v)

    def bfs(self, start):
        visited = set()
        queue = deque([start])

        while queue:
            node = queue.popleft()
            if node not in visited:
                print(node, end=" ")
                visited.add(node)
                queue.extend(neighbor for neighbor in self.graph[node] if neighbor not in visited)

if __name__ == "__main__":

    g = Graph()
    g.add_edge(0, 1)
    g.add_edge(0, 2)
    g.add_edge(1, 2)
    g.add_edge(2, 0)
    g.add_edge(2, 3)
    g.add_edge(3, 3)

    start_node = 0
    print(f"BFS starting from node {start_node}:")
    g.bfs(start_node)


import binascii
import struct
import bpy
import rhino3dm as r3d
from bpy_extras.node_shader_utils import PrincipledBSDFWrapper
from . import utils

### default Rhino material name
DEFAULT_RHINO_MATERIAL = "Rhino Default Material"

#### material hashing functions

_black = (0, 0, 0, 255)

def Bbytes(b):
    """
    Return bytes representation of boolean
    """
    return struct.pack("?", b)

def Fbytes(f):
    """
    Return bytes representation of float
    """
    return struct.pack("f", f)

def Cbytes(c):
    """
    Return bytes representation of Color, a 4-tuple containing integers
    """
    return struct.pack("IIII", *c)

def tobytes(d):
    t = type(d)
    if t is bool:
        return Bbytes(d)
    if t is float:
        return Fbytes(d)
    if t is tuple and len(d) == 4:
        return Cbytes(d)

def hash_color(C, crc):
    """
    return crc from color C
    """
    crc = binascii.crc32(tobytes(C), crc)
    return crc

def hash_material(M):
    """
    Hash a rhino3dm.Material. A CRC32 is calculated using the
    material name and data that affects render results
    """
    crc = 13
    crc = binascii.crc32(bytes(M.Name, "utf-8"))
    crc = hash_color(M.DiffuseColor, crc)
    crc = hash_color(M.EmissionColor, crc)
    crc = hash_color(M.ReflectionColor, crc)
    crc = hash_color(M.SpecularColor, crc)
    crc = hash_color(M.TransparentColor, crc)
    crc = binascii.crc32(tobytes(M.DisableLighting), crc)
    crc = binascii.crc32(tobytes(M.FresnelIndexOfRefraction), crc)
    crc = binascii.crc32(tobytes(M.FresnelReflections), crc)
    crc = binascii.crc32(tobytes(M.IndexOfRefraction), crc)
    crc = binascii.crc32(tobytes(M.ReflectionGlossiness), crc)
    crc = binascii.crc32(tobytes(M.Reflectivity), crc)
    crc = binascii.crc32(tobytes(M.RefractionGlossiness), crc)
    crc = binascii.crc32(tobytes(M.Shine), crc)
    crc = binascii.crc32(tobytes(M.Transparency), crc)
    return crc

def material_name(m):
    h = hash_material(m)
    return m.Name + "~" + str(h)

def handle_materials(context, model, materials, update):
    """
    """
    for m in model.Materials:
        matname = material_name(m)
        if matname not in materials:
            blmat = utils.get_iddata(context.blend_data.materials, None, m.Name, None)
            if update:
                blmat.use_nodes = True
                refl = m.Reflectivity
                transp = m.Transparency
                ior = m.IndexOfRefraction
                roughness = m.ReflectionGlossiness
                transrough = m.RefractionGlossiness
                spec = m.Shine / 255.0

                if m.DiffuseColor == _black and m.Reflectivity > 0.0 and m.Transparency == 0.0:
                    r, g, b, _ = m.ReflectionColor
                elif m.DiffuseColor == _black and m.Reflectivity == 0.0 and m.Transparency > 0.0:
                    r, g, b, _ = m.TransparentColor
                    refl = 0.0
                elif m.DiffuseColor == _black and m.Reflectivity > 0.0 and m.Transparency > 0.0:
                    r, g, b, _ = m.TransparentColor
                    refl = 0.0
                else:
                    r, g, b, a = m.DiffuseColor
                    if refl > 0.0 and transp > 0.0:
                        refl = 0.0
                principled = PrincipledBSDFWrapper(blmat, is_readonly=False)
                principled.base_color = (r/255.0, g/255.0, b/255.0)
                principled.metallic = refl
                principled.transmission = transp
                principled.ior = ior
                principled.roughness = roughness
                principled.specular = spec
                principled.node_principled_bsdf.inputs[16].default_value = transrough
            materials[matname] = blmat