from __future__ import absolute_import, division, print_function

import warnings
import tensorflow as tf
import os
import sys
import keras.backend as K
import numpy as np
os.environ["CUDA_VISIBLE_DEVICES"] = sys.argv[1]
tf.enable_eager_execution()
from keras.datasets import cifar10
import keras.callbacks as callbacks
import keras.utils.np_utils as kutils
from keras import regularizers

_BATCH_NORM_DECAY = 0.997
_BATCH_NORM_EPSILON = 1e-5
# DEFAULT_DTYPE = tf.float32

class Resnet(tf.keras.Model):

    def build_blocks(self, out_filters, num_blocks, stride):
        blocks = []
        strides = [stride] + [1]*(num_blocks-1)
        for _stride in strides:
            blocks.append(self.make_basic_block(self.in_filters, out_filters, _stride))
            self.in_filters = out_filters
        return blocks

    def make_basic_block(self, in_filters, out_filters, stride=1):

        block = []

        layers = tf.keras.Sequential([
            tf.keras.layers.ZeroPadding2D(1),
            tf.keras.layers.Conv2D(filters = out_filters, kernel_size = 3, strides= stride, padding = "valid", use_bias = False, kernel_initializer='VarianceScaling', kernel_regularizer=regularizers.l2(self.wd)),
            tf.keras.layers.BatchNormalization(axis=self.channel_axis, momentum=_BATCH_NORM_DECAY, epsilon=_BATCH_NORM_EPSILON, center=True,
                                                                            scale=True, fused=True),
            tf.keras.layers.Activation('relu'),

            tf.keras.layers.ZeroPadding2D(1),
            tf.keras.layers.Conv2D(filters = out_filters, kernel_size = 3, strides= 1, padding = "valid", use_bias = False, kernel_initializer='VarianceScaling', kernel_regularizer=regularizers.l2(self.wd)),
            tf.keras.layers.BatchNormalization(axis=self.channel_axis, momentum=_BATCH_NORM_DECAY, epsilon=_BATCH_NORM_EPSILON, center=True,
                                                                            scale=True, fused=True),

            # tf.keras.layers.Activation('relu')
            ])

        if stride != 1 or in_filters != out_filters:
            shortcut = tf.keras.Sequential([
                tf.keras.layers.Conv2D(filters = out_filters, kernel_size = 1, strides= stride, padding = "valid", use_bias = False, kernel_initializer='VarianceScaling', kernel_regularizer=regularizers.l2(self.wd)),
                 tf.keras.layers.BatchNormalization(axis=self.channel_axis, momentum=_BATCH_NORM_DECAY, epsilon=_BATCH_NORM_EPSILON, center=True,
                                                                            scale=True, fused=True)])
        else:
            shortcut = tf.keras.Sequential([])

        block.append(layers)
        block.append(shortcut)

        return block

    def _create_ResnetModel(self, filters):
        # filters : Number of fliters used in 1st blk (after initial layer)
        # kernel_size : kerel size for all the basic building blocks.
        # blocklist : List of values, each value shows the number of basic building blocks in the ith blk.
        # strides : stride for the 1st conv layer of all 4 blks (each basic building blk made of 2 conv layers)
        # For resnet18, blocklist = [2, 2, 2, 2], eack blk made of 2 basic building blocks.

        self.intial_conv = tf.keras.Sequential([
            tf.keras.layers.ZeroPadding2D(1),
            tf.keras.layers.Conv2D(filters=64, kernel_size=3, strides=1, padding = "valid", data_format = self.data_format, use_bias = False, kernel_initializer='VarianceScaling', kernel_regularizer=regularizers.l2(self.wd)),
            tf.keras.layers.BatchNormalization(axis=self.channel_axis, momentum=_BATCH_NORM_DECAY, epsilon=_BATCH_NORM_EPSILON, center=True,
                                                                            scale=True, fused=True),
            tf.keras.layers.Activation('relu')
            ])

        self.blocks_1 = self.build_blocks(64, self.block_list[0], stride=1)
        self.blocks_2 = self.build_blocks(128, self.block_list[1], stride=2)
        self.blocks_3 = self.build_blocks(256, self.block_list[2], stride=2)
        self.blocks_4 = self.build_blocks(512, self.block_list[3], stride=2)

        self.final_dense = tf.keras.Sequential([
            tf.keras.layers.AveragePooling2D(pool_size=4),
            tf.keras.layers.Flatten(),
            tf.keras.layers.Dense(self.classes, activation='softmax', kernel_regularizer=regularizers.l2(self.wd))
            ])

    # resnet with basic building blocks
    def __init__(self, data_format, initial_filters=64, block_list=[2, 2, 2, 2], classes=10, wt_decay = 0.001):

        """training: Either True or False, whether we are currently training the
           model. Needed for batch norm.
          data_format: The input format ('channels_last' or 'channels_first')."""

        super(Resnet, self).__init__()

        self.in_filters = initial_filters
        self.block_list = block_list
        self.num_blocks = len(block_list)
        self.data_format = data_format
        self.wd = wt_decay
        self.classes = classes
        # self.training = training
        self.channel_axis = 1 if self.data_format == 'channels_first' else -1

        self._create_ResnetModel(filters = initial_filters)

        # self.avg_pool = tf.keras.layers.GlobalAveragePooling2D()
        # self.flatten = tf.keras.layers.Flatten(data_format = self.data_format)
        # self.fc = tf.keras.layers.Dense(self.classes, kernel_regularizer=regularizers.l2(self.wd))

    def call(self, inputs):
        """if self.data_format == 'channels_first':
            # Convert the inputs from channels_last (NHWC) to channels_first (NCHW).
            # This provides a large performance boost on GPU. See
            # https://www.tensorflow.org/performance/performance_guide#data_formats
            inputs = tf.transpose(inputs, [0, 3, 1, 2])"""

        inputs = self.intial_conv(inputs)

        for block in self.blocks_1:
            block_x = block[0](inputs)
            block_y = block[1](inputs)
            inputs = tf.keras.layers.Activation('relu')(block_x+ block_y)

        for block in self.blocks_2:
            block_x = block[0](inputs)
            block_y = block[1](inputs)
            inputs = tf.keras.layers.Activation('relu')(block_x+ block_y)

        for block in self.blocks_3:
            block_x = block[0](inputs)
            block_y = block[1](inputs)
            inputs = tf.keras.layers.Activation('relu')(block_x+ block_y)

        for block in self.blocks_4:
            block_x = block[0](inputs)
            block_y = block[1](inputs)
            inputs = tf.keras.layers.Activation('relu')(block_x+ block_y)

        return self.final_dense(inputs)

if __name__ == '__main__':

    batch_size = 128
    img_rows, img_cols = 32, 32
    epochs = 200

    (trainX, trainY), (testX, testY) = cifar10.load_data()

    trainX = trainX.astype('float32')
    trainX = (trainX - trainX.mean(axis=0)) / (trainX.std(axis=0))
    testX = testX.astype('float32')
    testX = (testX - testX.mean(axis=0)) / (testX.std(axis=0))

    trainY = kutils.to_categorical(trainY, num_classes = 10)
    testY = kutils.to_categorical(testY, num_classes = 10)

    #testY = testY.astype(np.int64)
    #trainY = trainY.astype(np.int64)

    print(K.image_data_format())
    model = Resnet(data_format='channels_last')

    model._set_inputs(tf.zeros((batch_size, 32, 32, 3)))

    model.compile(optimizer=tf.train.AdamOptimizer(0.001), loss='categorical_crossentropy',
                      metrics=['accuracy'])

    model.fit(trainX, trainY, batch_size=batch_size, epochs=epochs,validation_data=(testX, testY), verbose=1)
    scores = model.evaluate(testX, testY, batch_size, verbose=1)


import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.utils as vutils
from torch.autograd import Variable
import itertools, time, os
from glob import glob
import models.cyclegan as cyclegan

def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        m.weight.data.normal_(0.0, 0.02)

    elif classname.find('BatchNorm') != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)

class Trainer(object):
    def __init__(self, config, dataloader):
        self.config = config
        self.a_data_loader = dataloader[0]
        self.b_data_loader = dataloader[1]

        self.ngpu = int(config.ngpu)
        self.input_nc = int(config.input_nc)
        self.output_nc = int(config.output_nc)

        self.ngf = int(config.ngf)
        self.ndf = int(config.ndf)
        self.cuda = config.cuda

        self.num_steps = len(self.a_data_loader)
        self.batch_size = config.batch_size
        self.image_size = config.image_size

        self.lr = config.lr
        self.beta1 = config.beta1
        self.beta2 = config.beta2

        self.niter = config.niter
        self.decay_epoch = config.decay_epoch
        self.cycle_lambda = config.cycle_lambda

        self.outf = config.outf
        self.sample_step = config.sample_step
        self.checkpoint_step = config.checkpoint_step

        self.build_model()

        if self.cuda:
            self.netG_AB.cuda()
            self.netD_A.cuda()
            self.netG_BA.cuda()
            self.netD_B.cuda()

    def load_model(self):
        print("[*] Load models from {}...".format(self.outf))

        paths = glob(os.path.join(self.outf, 'net*.pth'))
        paths.sort()

        if len(paths) == 0:
            print("[!] No checkpoint found in {}...".format(self.outf))
            return

        epochs = [int(os.path.basename(path.split('.')[0].split('_')[-2].split('-')[-1])) for path in paths]
        self.start_epoch = str(max(epochs))
        steps = [int(os.path.basename(path.split('.')[0].split('_')[-1].split('-')[-1])) for path in paths]
        self.start_step = str(max(steps))

        G_AB_filename = '{}/netG_A_epoch-{}_step-{}.pth'.format(self.outf, self.start_epoch, self.start_step)
        G_BA_filename = '{}/netG_B_epoch-{}_step-{}.pth'.format(self.outf, self.start_epoch, self.start_step)
        D_A_filename = '{}/netD_A_epoch-{}_step-{}.pth'.format(self.outf, self.start_epoch, self.start_step)
        D_B_filename = '{}/netD_B_epoch-{}_step-{}.pth'.format(self.outf, self.start_epoch, self.start_step)

        self.netG_AB.load_state_dict(torch.load(G_AB_filename))
        self.netG_BA.load_state_dict(torch.load(G_BA_filename))
        self.netD_A.load_state_dict(torch.load(D_A_filename))
        self.netD_B.load_state_dict(torch.load(D_B_filename))

        print("[*] Model loaded: {}".format(G_AB_filename))

    def build_model(self):
        self.netG_AB = cyclegan._netG(self.ngpu, self.ngf,
                                      self.input_nc, self.output_nc)
        self.netG_AB.apply(weights_init)

        self.netG_BA = cyclegan._netG(self.ngpu, self.ngf,
                                      self.input_nc, self.output_nc)
        self.netG_BA.apply(weights_init)

        self.netD_A = cyclegan._netD(self.ngpu, self.ndf, self.input_nc)
        self.netD_A.apply(weights_init)

        self.netD_B = cyclegan._netD(self.ngpu, self.ndf, self.input_nc)
        self.netD_B.apply(weights_init)

        if self.config.model_path != '':
            self.load_model()

    def train(self):
        MSELoss = nn.MSELoss()
        L1loss = nn.L1Loss()

        if self.cuda:
            MSELoss.cuda()
            L1loss.cuda()

        # setup optimizer
        optimizerD_A = optim.Adam(self.netD_A.parameters(), lr=self.lr, betas=(self.beta1, self.beta2))
        optimizerD_B = optim.Adam(self.netD_B.parameters(), lr=self.lr, betas=(self.beta1, self.beta2))
        optimizerG = optim.Adam(itertools.chain(self.netG_AB.parameters(), self.netG_BA.parameters()), lr=self.lr, betas=(self.beta1, self.beta2))

        A_loader, B_loader = iter(self.a_data_loader), iter(self.b_data_loader)
        valid_x_A, valid_x_B = self._get_variable(A_loader.next()), self._get_variable(B_loader.next())

        vutils.save_image(valid_x_A.data, '{}/valid_x_A.png'.format(self.outf), nrow=10)
        vutils.save_image(valid_x_B.data, '{}/valid_x_B.png'.format(self.outf), nrow=10)

        start_time = time.time()
        for epoch in range(self.niter):
            if (epoch+1) > self.decay_epoch:
                optimizerD_A.param_groups[0]['lr'] -= self.lr / (self.niter - self.decay_epoch)
                optimizerD_B.param_groups[0]['lr'] -= self.lr / (self.niter - self.decay_epoch)
                optimizerG.param_groups[0]['lr'] -= self.lr / (self.niter - self.decay_epoch)

            for step in range(self.num_steps):
                try:
                    realA, realB = A_loader.next(), B_loader.next()
                except StopIteration:
                    A_loader, B_loader = iter(self.a_data_loader), iter(self.b_data_loader)
                    realA, realB = A_loader.next(), B_loader.next()
                if realA.size(0) != realB.size(0):
                    print("[!] Sampled dataset from A and B have different # of data. Try resampling...")
                    continue

                batch_size = realA.size(0)
                realA, realB = Variable(realA.cuda()), Variable(realB.cuda())
                ############################
                # (1) Update G network: minimize Lgan(MSE) + Lcycle(L1)
                ###########################
                for p in self.netD_A.parameters():
                    p.requires_grad = False
                for p in self.netD_B.parameters():
                    p.requires_grad = False

                self.netG_AB.zero_grad()
                self.netG_BA.zero_grad()

                # GAN loss: D_B(G_A(A))
                fakeB = self.netG_AB(realA)
                output = self.netD_B(fakeB)
                loss_G_A = MSELoss(output, Variable(torch.ones(output.size()).cuda()))

                # GAN loss: D_A(G_B(B))
                fakeA = self.netG_BA(realB)
                output = self.netD_A(fakeA)
                loss_G_B = MSELoss(output, Variable(torch.ones(output.size()).cuda()))

                # Forward cycle loss: A <-> G_B(G_A(A))
                cycleA = self.netG_BA(fakeB)
                loss_cycle_A = L1loss(cycleA, realA) * self.cycle_lambda

                # Backward cycle loss: B <-> G_A(G_B(B))
                cycleB = self.netG_AB(fakeA)
                loss_cycle_B = L1loss(cycleB, realB) * self.cycle_lambda

                # Combined Generator loss
                loss_G = loss_G_A + loss_G_B + loss_cycle_A + loss_cycle_B
                loss_G.backward()
                optimizerG.step()

                ############################
                # (2) Update D network: minimize LSGAN loss
                ###########################
                for p in self.netD_A.parameters():
                    p.requires_grad = True
                for p in self.netD_B.parameters():
                    p.requires_grad = True

                ## train D_A ##
                # train with real
                self.netD_A.zero_grad()

                D_A_real = self.netD_A(realA)
                loss_D_A_real = MSELoss(D_A_real, Variable(torch.ones(D_A_real.size()).cuda()))

                # train with fake
                D_A_fake = self.netD_A(fakeA.detach())
                loss_D_A_fake = MSELoss(D_A_fake, Variable(torch.zeros(D_A_fake.size()).cuda()))

                loss_D_A = loss_D_A_real + loss_D_A_fake
                loss_D_A.backward()
                optimizerD_A.step()

                ## train D_B ##
                # train with real
                self.netD_B.zero_grad()

                D_B_real = self.netD_B(realB)
                loss_D_B_real = MSELoss(D_B_real, Variable(torch.ones(D_B_real.size()).cuda()))

                # train with fake
                D_B_fake = self.netD_B(fakeB.detach())
                loss_D_B_fake = MSELoss(D_B_fake, Variable(torch.zeros(D_B_fake.size()).cuda()))

                loss_D_B = loss_D_B_real + loss_D_B_fake
                loss_D_B.backward()
                optimizerD_B.step()

                step_end_time = time.time()

                print('[%d/%d][%d/%d] - time_passed: %.2f, loss_D_A: %.3f, loss_D_B: %.3f, '
                      'loss_G_A: %.3f, loss_G_B: %.3f, loss_A_cycle: %.3f, loss_B_cycle: %.3f'
                      % (epoch, self.niter, step, len(self.a_data_loader), step_end_time - start_time,
                         loss_D_A, loss_D_B, loss_G_A, loss_G_B, loss_cycle_A, loss_cycle_B))

                if step % self.sample_step == 0:
                    fake = self.netG_AB(valid_x_A)
                    vutils.save_image(fake.data, '%s/AB_%03d_%d.png' % (self.outf, epoch, step),
                                      nrow=10)
                    fake = self.netG_BA(valid_x_B)
                    vutils.save_image(fake.data, '%s/BA_%03d_%d.png' % (self.outf, epoch, step),
                                      nrow=10)

                if step% self.checkpoint_step == 0 and step != 0:
                    torch.save(self.netG_AB.state_dict(), '%s/netG_A_epoch-%d_step-%s.pth' % (self.outf, epoch, step))
                    torch.save(self.netD_A.state_dict(), '%s/netD_A_epoch-%d_step-%s.pth' % (self.outf, epoch, step))
                    torch.save(self.netG_BA.state_dict(), '%s/netG_B_epoch-%d_step-%s.pth' % (self.outf, epoch, step))
                    torch.save(self.netD_B.state_dict(), '%s/netD_B_epoch-%d_step-%s.pth' % (self.outf, epoch, step))
                    print("Saved checkpoint")

    def _get_variable(self, inputs):
        if self.ngpu > 0:
            out = Variable(inputs.cuda())
        else:
            out = Variable(inputs)
        return out


import torch
import torchvision.transforms.functional as TF
import torchvision.transforms as T
import kornia.augmentation as K
from open_clip import tokenize
import random
import src.templates as templates

def denormalize(image):
    return K.Denormalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))(
        image)

def normalize(image):
    return T.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))(
        image)

def generate_random_param(augmentation_type, include_org=False):
    if include_org:
        low = 0
    else:
        low = 1

    if augmentation_type in (
            'vflip', 'invert', 'grayscale', 'random_erase', 'elastic_transform', 'solarize',
            'bright', 'dark', 'blur'):
        param = torch.randint(low=low, high=2, size=(1,)).item()
    elif augmentation_type == 'rotation':
        param = torch.randint(low=low, high=4, size=(1,)).item()
    elif augmentation_type in ('snow', 'frost', 'fog', 'noise', 'jpeg'):
        param = torch.randint(low=low, high=6, size=(1,)).item()
    else:
        raise Exception(f'Unknown augmentation type {augmentation_type}')

    return param

def augment_image(augmentation_type, image, param=None):
    if param is None:
        param = generate_random_param(augmentation_type)

    if augmentation_type == 'vflip':
        if param == 0:
            image = image
        elif param == 1:
            image = TF.vflip(image)
    elif augmentation_type == 'rotation':
        if param == 0:
            image = image
        elif param == 1:
            image = TF.rotate(image, 90)
        elif param == 2:
            image = TF.rotate(image, 180)
        elif param == 3:
            image = TF.rotate(image, 270)
    elif augmentation_type == 'random_erase':
        if param == 0:
            image = image
        elif param == 1:
            image = T.RandomErasing(p=1)(image)
    elif augmentation_type == 'elastic_transform':
        if param == 0:
            image = image
        elif param == 1:
            image = T.ElasticTransform(alpha=100.)(image)
    elif augmentation_type == 'invert':
        if param == 0:
            image = image
        elif param == 1:
            image = normalize(TF.invert(denormalize(image)))
    elif augmentation_type == 'solarize':
        if param == 0:
            image = image
        elif param == 1:
            image = normalize(TF.solarize(denormalize(image), 0.6))
    elif augmentation_type == 'grayscale':
        if param == 0:
            image = image
        elif param == 1:
            image = normalize(TF.rgb_to_grayscale(denormalize(image), num_output_channels=3))
    elif augmentation_type == 'bright':
        if param == 0:
            image = image
        elif param == 1:
            image = normalize(TF.adjust_brightness(denormalize(image), brightness_factor=2))
    elif augmentation_type == 'dark':
        if param == 0:
            image = image
        elif param == 1:
            image = normalize(TF.adjust_brightness(denormalize(image), brightness_factor=0.5))
    elif augmentation_type == 'blur':
        if param == 0:
            image = image
        elif param == 1:
            image = TF.gaussian_blur(image, 17)
    elif augmentation_type == 'snow':
        return image
    elif augmentation_type == 'frost':
        return image
    elif augmentation_type == 'fog':
        return image
    elif augmentation_type == 'noise':
        return image
    elif augmentation_type == 'jpeg':
        return image

    else:
        raise Exception(f'Unknown augmentation type {augmentation_type}')

    return image


from django.shortcuts import render
from django.shortcuts import redirect
from django import forms
from . import util
import markdown2
import random
from . import util

class NewArticleForm(forms.Form):
    title = forms.CharField(widget=forms.TextInput(attrs={"class":"form-control", "id":"newTitle"}))
    body = forms.CharField(widget=forms.Textarea(attrs={"rows":"10", "class":"form-control", "id":"newBody"}))

    def clean_title(self):
        data = self.cleaned_data.get("title")
        entries = util.list_entries()

        if data.lower() in [entry.lower() for entry in entries]:
            raise forms.ValidationError("That article already exists")

        return data

class EditArticle(forms.Form):
    title = forms.CharField(widget=forms.TextInput(
        attrs={"class": "form-control", "id": "newTitle"}))
    body = forms.CharField(widget=forms.Textarea(
        attrs={"rows": "10", "class": "form-control", "id": "newBody"}))

    def clean_body(self):
        return self.cleaned_data.get("body")

def index(request):
    return render(request, "encyclopedia/index.html", {
        "entries": util.list_entries()
    })

def lookup(request, entry):
    entry_content = util.get_entry(entry)
    if entry_content:
        entry_content = markdown2.markdown(entry_content)
    else:
        entry_content = f"
Error: there is no page for \"{entry}\"
"
    return render(request, "encyclopedia/entry.html", {
        "entry_name": entry,
        "entry_content": entry_content
    })

def random_page(request):
    entry = random.choice(util.list_entries())
    return redirect("encyclopedia:lookup", entry=entry)

def search(request):
    query = request.GET['q']
    entries = util.list_entries()
    if query.lower() in [lower_entry.lower() for lower_entry in entries]:
        return redirect("encyclopedia:lookup", entry=query)
    else:
        search_results = [entry for entry in entries if query.lower() in entry.lower()]
        return render(request, "encyclopedia/search.html", {
            "query":query,
            "search_results":search_results
        })

def new(request):
    if request.method == "POST":
        form = NewArticleForm(request.POST)
        if form.is_valid():
            form_title = form.cleaned_data["title"]
            with open(f"entries/{form_title}.md", "w") as file:
                file.write(form.cleaned_data["body"])
            return redirect("encyclopedia:lookup", entry=form_title)
        else:
            return render(request, "encyclopedia/new.html", {
                "form": form
            })
    else:
        return render(request, "encyclopedia/new.html", {
            "form": NewArticleForm()
        })

def edit(request, entry):
    if request.method == "GET":
        form = EditArticle(initial={"title": entry, "body": util.get_entry(entry)})

        return render(request, "encyclopedia/edit.html", {
            "form": form,
            "entry": entry
        })
    else:
        form = EditArticle(request.POST)
        if form.is_valid():
            form_title = form.cleaned_data["title"]
            util.save_entry(form_title, form.cleaned_data["body"])
            return redirect("encyclopedia:lookup", entry=form_title)
        else:
            return render(request, "encyclopedia/edit.html", {
                "form": form,
                "entry": entry
            })


# Line editor

import uibase

def edit_line(line: str, line_number: int):
    uibase.cls()
    print('Editing: {} line\n'.format(line_number))
    uibase.simulate_input(line)
    new_line = input().rstrip()
    return new_line

def replace_line(line: str, line_number: int):
    uibase.cls()
    print('Replacing: {} line\n'.format(line_number))
    new_line = input().rstrip()
    return new_line

def free_input(line_number: int):
    uibase.cls()
    print('Inserting lines to: {}line\n'.format(line_number))
    lines = []
    while True:
        try:
            lines.append(input().rstrip())
        except KeyboardInterrupt:
            break
    return lines