import pickle as pkl
import torch
import torch.nn as nn
from torch.autograd import Variable
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

import numpy as np
import matplotlib.pyplot as plt

# Hyper-parameters
input_size = 784
hidden_size = 500
num_classes = 10
num_epochs = 5
batch_size = 100
learning_rate = 0.001

file_name_og = "Layer_surface_code\\cnn_mnist_models\\model-original-2.ckpt"
file_name_rob = "Layer_surface_code\\cnn_mnist_models\\model-robust_2.ckpt"

# Fully connected neural network with one hidden layer
class NeuralNet(nn.Module):
    def __init__(self):
        super(NeuralNet, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.layer2 = nn.Sequential(
            nn.Conv2d(32, 64, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.drop_out = nn.Dropout()
        self.fc1 = nn.Linear(7 * 7 * 64, 1000)
        self.fc2 = nn.Linear(1000, 10)

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.reshape(out.size(0), -1)
        out = self.drop_out(out)
        out = self.fc1(out)
        out = self.fc2(out)

        return out

def train(device, train_loader):
    model = NeuralNet().to(device)

    # Loss and optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

    # Train the model
    total_step = len(train_loader)
    for epoch in range(num_epochs):
        for i, (images, labels) in enumerate(train_loader):
            # Move tensors to the configured device
            images = images.to(device)
            labels = labels.to(device)

            # Forward pass
            outputs = model(images)
            loss = criterion(outputs, labels)

            # Backward and optimize
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            if (i + 1) % 100 == 0:
                print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'
                      .format(epoch + 1, num_epochs, i + 1, total_step, loss.item()))

        # Save the model checkpoint
        torch.save(model.state_dict(), 'Layer_surface_code/cnn_mnist_models/model-{}.ckpt'.format(epoch))

    return model

def test(device, test_loader, eps, file_name):
    # Test the model
    # In test phase, we don't need to compute gradients (for memory efficiency)
    criterion = nn.CrossEntropyLoss()

    model = NeuralNet().to(device)
    model.load_state_dict(torch.load(file_name))

    correct = 0
    adv_correct = 0
    misclassified = 0
    total = 0
    noises = []
    y_preds = []
    y_preds_adv = []

    for images, labels in test_loader:
        images = Variable(images.to(device), requires_grad=True)
        labels = Variable(labels.to(device))

        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()

        # Add perturbation
        grad = torch.sign(images.grad.data)
        imgs_adv = torch.clamp(images.data + (eps * grad), 0, 1)

        adv_outputs = model(Variable(imgs_adv))

        _, predicted = torch.max(outputs.data, 1)
        _, adv_preds = torch.max(adv_outputs.data, 1)

        total += labels.size(0)
        correct += (predicted == labels).sum().item()
        adv_correct += (adv_preds == labels).sum().item()
        misclassified += (predicted != adv_preds).sum().item()

        noises.extend((images - imgs_adv).data.numpy())
        y_preds.extend(predicted.data.numpy())
        y_preds_adv.extend(adv_preds.data.numpy())

    print('Accuracy of the network w/o adversarial attack on the 10000 test images: {} %'.format(100 * correct / total))
    print('Accuracy of the network with adversarial attack on the 10000 test images: {} %'.format(
        100 * adv_correct / total))
    print('Number of misclassified examples (as compared to clean predictions): {}/{}'.format(misclassified, total))
    print("\n")

    return 100*(correct/total), 100*(adv_correct/total), (misclassified, total)

def main(flag):
    # Device configuration
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # MNIST dataset
    train_dataset = torchvision.datasets.MNIST(root='./data',
                                               train=True,
                                               transform=transforms.ToTensor(),
                                               download=True)

    test_dataset = torchvision.datasets.MNIST(root='./data',
                                              train=False,
                                              transform=transforms.ToTensor())

    # Data loader
    train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)

    test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)

    if flag == "train":
        model = train(device, train_loader)
    elif flag == "test":
        pertub = np.linspace(0.001, 0.3, 10)
        no_attack_og = []
        attack_og = []
        no_attack_rob = []
        attack_rob = []

        for i in range(len(pertub)):
            og = test(device, test_loader, pertub[i], file_name=file_name_og)
            rob = test(device, test_loader, pertub[i], file_name=file_name_rob)

            no_attack_og.append(og[0])
            attack_og.append(og[1])

            no_attack_rob.append(rob[0])
            attack_rob.append(rob[1])

        plt.figure(1)
        plt.plot(pertub, attack_og, 'r*')
        plt.plot(pertub, attack_og, 'r', label='Original Model')

        # plt.plot(pertub, no_attack_og, 'g*')
        # plt.plot(pertub, no_attack_og, 'g', label='Original Model')

        plt.plot(pertub, attack_rob, 'b*')
        plt.plot(pertub, attack_rob, 'b', label='Robust Model')

        # plt.plot(pertub, no_attack_rob, 'y', label='Robust Model')

        plt.ylabel('Model Accuracy (%)')
        plt.xlabel('Perturbation magnitude')
        plt.legend()
        plt.show()

if __name__ == "__main__":
    main("test")


from random import randint

def main():
    n = 10000
    count = 0
    wins = 0

    for i in range(n):
        choice = getVal()
        car = getVal()

        opened = getVal(choice, car)

        if opened == car:
            continue

        count += 1

        if choice != car:
            wins += 1

    print(f"probability is: {wins / count:.3f}")

def getVal(choice=None, car=None):
    if choice is None:
        return randint(1,3)

    while True:
        temp = randint(1,3)

        if temp == choice or temp == car:
            continue

        return temp

if __name__ == '__main__':
    main()


# More details can be found in TechToTinker.blogspot.com
# George Bantique | tech.to.tinker@gmail.com

from mfrc522 import MFRC522
from machine import Pin
from machine import SPI

spi = SPI(2, baudrate=2500000, polarity=0, phase=0)
# Using Hardware SPI pins:
#     sck=18   # yellow
#     mosi=23  # orange
#     miso=19  # blue
#     rst=4    # white
#     cs=5     # green, DS
# *************************
# To use SoftSPI,
# from machine import SOftSPI
# spi = SoftSPI(baudrate=100000, polarity=0, phase=0, sck=sck, mosi=mosi, miso=miso)
spi.init()
rdr = MFRC522(spi=spi, gpioRst=4, gpioCs=5)
print("Place card")

while True:

    (stat, tag_type) = rdr.request(rdr.REQIDL)
    if stat == rdr.OK:
        (stat, raw_uid) = rdr.anticoll()
        if stat == rdr.OK:
            card_id = "uid: 0x%02x%02x%02x%02x" % (raw_uid[0], raw_uid[1], raw_uid[2], raw_uid[3])
            print(card_id)


import zmq
context = zmq.Context()
socket = context.socket(zmq.PULL)
socket.bind("tcp://*:6000")
poller = zmq.Poller()
poller.register(socket, zmq.POLLIN)
print('Looking out for messages!')
while True:
    socks = dict(poller.poll())
    if socket in socks:
         message = socket.recv_json()
         print(message)
         print("Received!!!!")


from collections import defaultdict
import gzip
import random
import math
import scipy
from scipy import optimize
import numpy
import nltk
from nltk.tokenize import word_tokenize
import string
from nltk.stem.porter import *
from sklearn import linear_model
from scipy import spatial
from random import shuffle
from nltk.corpus import stopwords
import itertools
from sklearn.metrics import mean_squared_error
import statistics

# read data
data = []
requiredEntrySet = {"fit", "bust size", "weight", "rating", "rented for", "review_text", "body type", "review_summary", "category", "height", "size", "age"}
def readJSON(path):
	for l in gzip.open(path, 'rt'):
		try:
			d = eval(l)
			curEntries = d.keys()
			if requiredEntrySet.issubset(curEntries):
				yield d
		except Exception as e:
			pass

for d in readJSON("renttherunway_final_data.json.gz"):
	data.append(d)

data = data[:10000]

# process the data to parse the height and the weight
import re

def parse_height(height):
    pattern = re.compile(r'(\d+)\' (\d+)\"')
    match = pattern.match(height)

    if match:
        feet = int(match.group(1))
        inch = int(match.group(2))

        return feet*12 + inch
    else:
        return -1

def parse_weight(weight):
    pattern = re.compile(r'(\d+)lbs')
    match = pattern.match(weight)
    if match:
        return int(match.group(1))
    else:
        return -1

def parse_data(data):
    for d in data:
        height = parse_height(d['height'])
        d['height'] = height

        weight = parse_weight(d['weight'])
        d['weight'] = weight

parse_data(data)

# extract the min and max for regulization
def extract_max(data):
    maxAge = -1
    maxSize = -1
    maxWeight = -1
    maxHeight = -1

    for d in data:
        if int(d["age"]) > maxAge:
            maxAge = int(d["age"])
        if d["size"] > maxSize:
            maxSize = d["size"]
        if d["weight"] > maxWeight:
            maxWeight = d["weight"]
        if d["height"] > maxHeight:
            maxHeight = d["height"]

    return maxAge, maxSize, maxWeight, maxHeight

maxAge, maxSize, maxWeight, maxHeight = extract_max(data)

def extract_min(data):
    minAge = maxAge
    minSize = maxSize
    minWeight = maxWeight
    minHeight = maxHeight

    for d in data:
        if int(d["age"]) < minAge:
            minAge = int(d["age"])
        if d["size"] < minSize:
            minSize = d["size"]
        if d["weight"] < minWeight:
            minWeight = d["weight"]
        if d["height"] < minHeight:
            minHeight = d["height"]

    return minAge, minSize, minWeight, minHeight

minAge, minSize, minWeight, minHeight = extract_min(data)

# extract the regulized features
def regulize(d, minD, maxD):
    return (d - minD)/(maxD - minD)

def age(datum, minAge, maxAge):
    return regulize(int(datum["age"]), minAge, maxAge)

def size(datum, minSize, maxSize):
    return regulize(datum["size"], minSize, maxSize)

def weight(datum, minWeight, maxWeight):
    return regulize(datum["weight"], minWeight, maxWeight)

def height(datum, minHeight, maxHeight):
    return regulize(datum["height"], minHeight, maxHeight)

# One hot encoding features
Entry = ["fit", "rented for", "body type", "category", "bust size"]

def extract_onehot(data, Entry):
	Maps = {Entry[i]:defaultdict(int) for i in range(len(Entry))}
	for d in data:
		for f in Entry:
			mp = Maps[f]
			if d[f] not in mp:
				mp[d[f]] = len(mp)
	return Maps

# fit: 3; rented for: 9; body type: 7; category: 68; bust size: 101

Maps = extract_onehot(data, Entry)

# Combined single vector
def feature_onehot_combined(datum):
	featureVectors = [[0]*len(Maps[f]) for f in Entry]
	for i in range(len(Entry)):
		f = Entry[i]
		featureVectors[i][Maps[f][d[f]]] = 1
	singleVector = list(itertools.chain.from_iterable(featureVectors))
	singleVector.append(1)
	return singleVector

# separate features (all 5)
def feature_onehot_separated(datum):
	featureVectors = [[0]*len(Maps[f]) for f in Entry]
	for i in range(len(Entry)):
		f = Entry[i]
		featureVectors[i][Maps[f][d[f]]] = 1
	for v in featureVectors:
		v.append(1)
	return featureVectors

# print(feature_onehot_separated(data[22]))

# fit feature
def featureFit(datum):
	feat = [0] * len(Maps['fit'])
	feat[Maps['fit'][d['fit']]] = 1
	feat.append(1)
	return feat

def featureRentedFor(datum):
	feat = [0] * len(Maps['rented for'])
	feat[Maps['rented for'][d['rented for']]] = 1
	feat.append(1)
	return feat

def featureBodyType(datum):
	feat = [0] * len(Maps['body type'])
	feat[Maps['body type'][d['body type']]] = 1
	feat.append(1)
	return feat

def featureCategory(datum):
	feat = [0] * len(Maps['category'])
	feat[Maps['category'][d['category']]] = 1
	feat.append(1)
	return feat

def featureBustSize(datum):
	feat = [0] * len(Maps['bust size'])
	feat[Maps['bust size'][d['bust size']]] = 1
	feat.append(1)
	return feat

#import nltk
#nltk.download('stopwords')

# unigram and bigram
uniWordCount = defaultdict(int)
biWordCount = defaultdict(int)

uniWordCount_rmPunc = defaultdict(int)
biWordCount_rmPunc = defaultdict(int)

uniWordCount_rmPunc_Stop = defaultdict(int)
biWordCount_rmPunc_Stop = defaultdict(int)

punctuation = set(string.punctuation)
stopwords = set(stopwords.words('english'))
punctuation.add('\n')
stemmer = PorterStemmer()

for d in data:
	text = ''.join([c for c in d['review_summary'].lower()])
	for token in text.split():
		uniWordCount[token] += 1
	text = ''.join([c for c in d['review_summary'].lower() if c not in punctuation ])
	for token in text.split():
		uniWordCount_rmPunc[token] += 1
	text = ''.join([c for c in d['review_summary'].lower() if c not in punctuation and c not in stopwords])
	for token in text.split():
		uniWordCount_rmPunc_Stop[token] += 1

	text = ''.join([c for c in d['review_text'].lower() if c not in punctuation ])
	bigrm = nltk.bigrams(nltk.word_tokenize(text))
	for token in bigrm:
		biWordCount_rmPunc[token] += 1
	text = ''.join([c for c in d['review_text'].lower()])
	bigrm = nltk.bigrams(nltk.word_tokenize(text))
	for token in bigrm:
		biWordCount[token] += 1
	text = ''.join([c for c in d['review_text'].lower() if c not in punctuation and c not in stopwords])
	bigrm = nltk.bigrams(nltk.word_tokenize(text))
	for token in bigrm:
		biWordCount_rmPunc_Stop[token] += 1

top_k_uni = 40
top_k_bi = 40

unicounts = [(uniWordCount[w], w) for w in uniWordCount]
unicounts.sort()
unicounts.reverse()
uniwords = [x[1] for x in unicounts[:top_k_uni]]
uniwordId = dict(zip(uniwords, range(len(uniwords))))
uniwordSet = set(uniwords)

unicounts_rmPunc = [(uniWordCount_rmPunc[w], w) for w in uniWordCount_rmPunc]
unicounts_rmPunc.sort()
unicounts_rmPunc.reverse()
uniwords_rmPunc = [x[1] for x in unicounts_rmPunc[:top_k_uni]]
uniwordId_rmPunc = dict(zip(uniwords_rmPunc, range(len(uniwords_rmPunc))))
uniwordSet_rmPunc = set(uniwords_rmPunc)

uniWordCount_rmPunc_Stop = [(uniWordCount_rmPunc_Stop[w], w) for w in uniWordCount]
uniWordCount_rmPunc_Stop.sort()
uniWordCount_rmPunc_Stop.reverse()
uniwords_rmPunc_Stop = [x[1] for x in uniWordCount_rmPunc_Stop[:top_k_uni]]
uniwordId_rmPunc_Stop = dict(zip(uniwords_rmPunc_Stop, range(len(uniwords_rmPunc_Stop))))
uniwordSet_rmPunc_Stop = set(uniwords_rmPunc_Stop)

bicounts = [(biWordCount[w], w) for w in biWordCount]
bicounts.sort()
bicounts.reverse()
biwords = [x[1] for x in bicounts[:top_k_bi]]
biwordId = dict(zip(biwords, range(len(biwords))))
biwordSet = set(biwords)

bicounts_rmPunc = [(biWordCount_rmPunc[w], w) for w in biWordCount_rmPunc]
bicounts_rmPunc.sort()
bicounts_rmPunc.reverse()
biwords_rmPunc = [x[1] for x in bicounts_rmPunc[:top_k_bi]]
biwordId_rmPunc = dict(zip(biwords_rmPunc, range(len(biwords_rmPunc))))
biwordSet_rmPunc = set(biwords_rmPunc)

bicounts_rmPunc_Stop = [(biWordCount_rmPunc_Stop[w], w) for w in biWordCount_rmPunc_Stop]
bicounts_rmPunc_Stop.sort()
bicounts_rmPunc_Stop.reverse()
biwords_rmPunc_Stop = [x[1] for x in bicounts_rmPunc_Stop[:top_k_bi]]
biwordId_rmPunc_Stop = dict(zip(biwords_rmPunc_Stop, range(len(biwords_rmPunc_Stop))))
biwordSet_rmPunc_Stop = set(biwords_rmPunc_Stop)

def featureSummary(datum, rmPunc, rmPuncStop):
	feat = [0]*len(uniwords)
	if rmPunc:
		r = ''.join([c for c in datum['review_summary'].lower() if c not in punctuation])
	elif rmPuncStop:
		r = ''.join([c for c in datum['review_summary'].lower() if c not in punctuation and c not in stopwords])
	else:
		r = ''.join([c for c in datum['review_summary'].lower()])
	for token in r.split():
		if rmPunc:
			if token in uniwords_rmPunc:
				feat[uniwordId_rmPunc[token]] += 1
		elif rmPuncStop:
			if token in uniwords_rmPunc_Stop:
				feat[uniwordId_rmPunc_Stop[token]] += 1
		else:
			if token in uniwords:
				feat[uniwordId[token]] += 1
	feat.append(1)
	return feat

def featureText(datum, rmPunc, rmPuncStop):
	feat = [0]*len(uniwords)
	if rmPunc:
		r = ''.join([c for c in datum['review_text'].lower() if c not in punctuation])
	elif rmPuncStop:
		r = ''.join([c for c in datum['review_text'].lower() if c not in punctuation and c not in stopwords])
	else:
		r = ''.join([c for c in datum['review_text'].lower()])

	for token in r.split():
		if rmPunc:
			if token in biwords_rmPunc:
				feat[biwordId_rmPunc[token]] += 1
		elif rmPuncStop:
			if token in biwords_rmPunc_Stop:
				feat[biwordId_rmPunc_Stop[token]] += 1
		else:
			if token in uniwords:
				feat[biwordId[token]] += 1
	feat.append(1)
	return feat

clf = linear_model.Ridge(1.0, fit_intercept=False) # MSE + 1.0 l2
clf.fit(X, y)
predictions = clf.predict(X_test)
print("mse of prediction is "+ str(mean_squared_error(y_test ,predictions)))


# Python script
#
# Author: Tetsuya Ishikawa <tiskw111@gmail.com>
# Date  : February 15, 2020
##################################################### SOURCE START #####################################################

import os
import pickle
import subprocess
import hashlib
import tarfile
import numpy as np

### Download a file from the given URL and check SHA256.
### This function returns True if and only if the SHA256 of the downloaded file is same as the ground truth.
def download_cifar10(filepath):

    url_origin = "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz"
    sha_256_gt = "6d958be074577803d12ecdefd02955f39262c83c16fe9348329d7fe0b5c001ce"

    ### Calculate SHA256 of the given file and compare with ground truth.
    def has_correct_sha256sum(filepath, hash_sha256_gt):
        with open(filepath, "rb") as ifp:
            data = ifp.read()
        hash_sha256 = hashlib.sha256(data).hexdigest()
        return hash_sha256 == hash_sha256_gt

    command = "wget -c %s -O %s" % (url_origin, filepath)
    subprocess.run(command, shell = True)
    return has_correct_sha256sum(filepath, sha_256_gt)

### Read file and unpickle from Tar and TarInfo data.
def read_and_unpickle(tar, tarinfo):

    return pickle.loads(tar.extractfile(tarinfo).read(), encoding = "bytes")

### This function extract all images from one batch data of CIFAR10.
### Each batch data contains 10,000 images and labels.
def parse_cifar_images(batch):

    Xs_batch = np.zeros((10000, 32, 32, 3))
    Xs_batch[:,:,:, 0] = batch[b"data"][:, 2048:    ].reshape((10000, 32, 32))
    Xs_batch[:,:,:, 1] = batch[b"data"][:, 1024:2048].reshape((10000, 32, 32))
    Xs_batch[:,:,:, 2] = batch[b"data"][:,    0:1024].reshape((10000, 32, 32))
    return Xs_batch

### This function extract all labels from one batch data of CIFAR10.
### Each batch data contains 10,000 images and labels.
def parse_cifar_labels(batch):

    return np.array(batch[b"labels"], dtype = np.uint8)

### Load and return CIFAR10 dataset as numpy array which is compatible with OpenCV image
### that has shape (height, width, channel) and the channel order is Blue-Green-Red.
### Returned values of this function is ((Xs_train, ys_train), (Xs_valid, ys_valid)), where
###   - Xs_train: (50000, 32, 32, 3),   - ys_train: (50000, ),
###   - Xs_valid: (10000, 32, 32, 3),   - ys_valid: (10000, ).
### If archive of the CIFAR10 is not exists, this function will download the archive from official website automatically.
def create_dataset(filepath):

    ### Download the CIFAR10 archive from the official website.
    if not os.path.exists(filepath):

        success_dl = download_cifar10(filepath)

        ### Abort if failed to download.
        if not success_dl:
            print("Hash (SHA256) does not match. You should remove %s and try again." % filepath)
            exit("Error: cifar10.py: Download failed (hash does not match)")

    ### Extract data and label files from the archive. The archive has the following files:
    ###   * cifar-10-batches-py/readme.html   (Readme file)
    ###   * cifar-10-batches-py/batches.meta  (Binary header)
    ###   * cifar-10-batches-py/data_batch_1  (Pickled images and labels for training)
    ###   * cifar-10-batches-py/data_batch_2  (Pickled images and labels for training)
    ###   * cifar-10-batches-py/data_batch_3  (Pickled images and labels for training)
    ###   * cifar-10-batches-py/data_batch_4  (Pickled images and labels for training)
    ###   * cifar-10-batches-py/data_batch_5  (Pickled images and labels for training)
    ###   * cifar-10-batches-py/test_batch    (Pickled images and labels for testing)
    ### and the following code extract only data files, un-pickle and store them as dictionary format.
    with tarfile.open(filepath, "r") as tar:
        data = {tarinfo.name:read_and_unpickle(tar, tarinfo) for tarinfo in tar if tarinfo.isreg() and "_batch" in tarinfo.name}

    ### Create training data.
    Xs_train = np.concatenate([parse_cifar_images(data["cifar-10-batches-py/data_batch_%d" % n]) for n in range(1, 6)])
    ys_train = np.concatenate([parse_cifar_labels(data["cifar-10-batches-py/data_batch_%d" % n]) for n in range(1, 6)])

    ### Create testing data.
    Xs_valid = parse_cifar_images(data["cifar-10-batches-py/test_batch"])
    ys_valid = parse_cifar_labels(data["cifar-10-batches-py/test_batch"])

    np.save("CIFAR10_train_images.npy", Xs_train)
    np.save("CIFAR10_train_labels.npy", ys_train)
    np.save("CIFAR10_test_images.npy",  Xs_valid)
    np.save("CIFAR10_test_labels.npy",  ys_valid)

if __name__  == "__main__":

    ### Create file path of the local cache of the CIFAR10 archive.
    current_dir = os.path.dirname(__file__)
    filepath    = os.path.join(current_dir, "cifar-10-python.tar.gz")

    create_dataset(filepath)

##################################################### SOURCE FINISH ####################################################
# vim: expandtab tabstop=4 shiftwidth=4 fdm=marker


import torch
import torch.nn as nn
from torch.nn import init
import torch.nn.functional as F
from torch.optim import lr_scheduler
import torchvision
from torchvision.models import resnet34

import functools
from einops import rearrange

import numpy as np

import models
from models.help_funcs import Transformer, TransformerDecoder, TwoLayerConv2d
from torch.nn.modules.padding import ReplicationPad2d
from models.ChangeFormer import ChangeFormerV1, ChangeFormerV6
###############################################################################
# Helper Functions
###############################################################################

def get_scheduler(optimizer, args):
    """Return a learning rate scheduler

    Parameters:
        optimizer          -- the optimizer of the network
        args (option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions．
                              opt.lr_policy is the name of learning rate policy: linear | step | plateau | cosine

    For 'linear', we keep the same learning rate for the first <opt.niter> epochs
    and linearly decay the rate to zero over the next <opt.niter_decay> epochs.
    For other schedulers (step, plateau, and cosine), we use the default PyTorch schedulers.
    See https://pytorch.org/docs/stable/optim.html for more details.
    """
    if args.lr_policy == 'linear':
        def lambda_rule(epoch):
            lr_l = 1.0 - epoch / float(args.max_epochs + 1)
            return lr_l
        scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
    elif args.lr_policy == 'step':
        step_size = args.max_epochs//3
        # args.lr_decay_iters
        scheduler = lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=0.1)
    elif args.lr_policy == 'multistep':
        scheduler = lr_scheduler.MultiStepLR(optimizer,
        milestones=[2, 4, 7, 11, 15, 25, 35, 47, 60, 70, 90, 110, 130, 150, 170, 180, 190], gamma=0.5)
    else:
        return NotImplementedError('learning rate policy [%s] is not implemented', args.lr_policy)
    return scheduler

class Identity(nn.Module):
    def forward(self, x):
        return x

def get_norm_layer(norm_type='instance'):
    """Return a normalization layer

    Parameters:
        norm_type (str) -- the name of the normalization layer: batch | instance | none

    For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev).
    For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics.
    """
    if norm_type == 'batch':
        norm_layer = functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True)
    elif norm_type == 'instance':
        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
    elif norm_type == 'none':
        norm_layer = lambda x: Identity()
    else:
        raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
    return norm_layer

def init_weights(net, init_type='normal', init_gain=0.02):
    """Initialize network weights.

    Parameters:
        net (network)   -- network to be initialized
        init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal
        init_gain (float)    -- scaling factor for normal, xavier and orthogonal.

    We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might
    work better for some applications. Feel free to try yourself.
    """
    def init_func(m):  # define the initialization function
        classname = m.__class__.__name__
        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
            if init_type == 'normal':
                init.normal_(m.weight.data, 0.0, init_gain)
            elif init_type == 'xavier':
                init.xavier_normal_(m.weight.data, gain=init_gain)
            elif init_type == 'kaiming':
                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
            elif init_type == 'orthogonal':
                init.orthogonal_(m.weight.data, gain=init_gain)
            else:
                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
            if hasattr(m, 'bias') and m.bias is not None:
                init.constant_(m.bias.data, 0.0)
        elif classname.find('BatchNorm2d') != -1:  # BatchNorm Layer's weight is not a matrix; only normal distribution applies.
            init.normal_(m.weight.data, 1.0, init_gain)
            init.constant_(m.bias.data, 0.0)

    print('initialize network with %s' % init_type)
    net.apply(init_func)  # apply the initialization function <init_func>

def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]):
    """Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights
    Parameters:
        net (network)      -- the network to be initialized
        init_type (str)    -- the name of an initialization method: normal | xavier | kaiming | orthogonal
        gain (float)       -- scaling factor for normal, xavier and orthogonal.
        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2

    Return an initialized network.
    """
    if len(gpu_ids) > 0:
        assert(torch.cuda.is_available())
        net.to(gpu_ids[0])
        if len(gpu_ids) > 1:
            net = torch.nn.DataParallel(net, gpu_ids)  # multi-GPUs
    init_weights(net, init_type, init_gain=init_gain)
    return net

def define_G(args, init_type='normal', init_gain=0.02, gpu_ids=[]):
    if args.net_G == 'base_resnet18':
        net = ResNet(input_nc=3, output_nc=2, output_sigmoid=False)

    elif args.net_G == 'base_transformer_pos_s4':
        net = BASE_Transformer(input_nc=3, output_nc=2, token_len=4, resnet_stages_num=4,
                             with_pos='learned')

    elif args.net_G == 'base_transformer_pos_s4_dd8':
        net = BASE_Transformer(input_nc=3, output_nc=2, token_len=4, resnet_stages_num=4,
                             with_pos='learned', enc_depth=1, dec_depth=8)
    elif args.net_G == 'base_transformer_pos_s4_dd8_o5':
        net = BASE_Transformer(input_nc=3, output_nc=5, token_len=4, resnet_stages_num=4,
                             with_pos='learned', enc_depth=1, dec_depth=8)

    elif args.net_G == 'base_transformer_pos_s4_dd8_dedim8':
        net = BASE_Transformer(input_nc=3, output_nc=2, token_len=4, resnet_stages_num=4,
                             with_pos='learned', enc_depth=1, dec_depth=8, decoder_dim_head=8)
    elif args.net_G == 'base_transformer_pos_s4_dd8_t8_e2d4':
       net = BASE_Transformer(input_nc=3, output_nc=2, token_len=8, resnet_stages_num=4,
                              with_pos='learned', enc_depth=2, dec_depth=4, decoder_dim_head=8)
    elif args.net_G == 'unet_coupled_trans_256':
       net = UNet_Change_Transformer()
    elif args.net_G == 'unet_coupled_two_trans_256':
       net = UNet_Change_Two_Transformer()
    elif args.net_G == 'siamUnet_conc':
       net = SiamUnet_conc(input_nbr=3, label_nbr=2)
    elif args.net_G == 'siamUnet':
       net = Res34_Unet_Double()
    elif args.net_G == 'changeFormer':
       net = ChangeFormerV1()
    elif args.net_G == 'changeFormerV6':
       net = ChangeFormerV6()
    elif args.net_G == 'newUNetTrans':
       net =  BASE_Transformer_UNet(input_nc=3, output_nc=2, token_len=4, resnet_stages_num=4,
                              with_pos='learned', with_decoder_pos='learned', enc_depth=1, dec_depth=8)
    else:
        raise NotImplementedError('Generator model name [%s] is not recognized' % args.net_G)
    return init_net(net, init_type, init_gain, gpu_ids)

###############################################################################
# main Functions
###############################################################################

class ResNet(torch.nn.Module):
    def __init__(self, input_nc, output_nc,
                 resnet_stages_num=5, backbone='resnet18',
                 output_sigmoid=False, if_upsample_2x=True):
        """
        In the constructor we instantiate two nn.Linear modules and assign them as
        member variables.
        """
        super(ResNet, self).__init__()
        expand = 1
        if backbone == 'resnet18':
            self.resnet = models.resnet18(pretrained=True,
                                          replace_stride_with_dilation=[False,True,True])
        elif backbone == 'resnet34':
            self.resnet = models.resnet34(pretrained=True,
                                          replace_stride_with_dilation=[False,True,True])
        elif backbone == 'resnet50':
            self.resnet = models.resnet50(pretrained=True,
                                          replace_stride_with_dilation=[False,True,True])
            expand = 4
        else:
            raise NotImplementedError
        self.relu = nn.ReLU()
        self.upsamplex2 = nn.Upsample(scale_factor=2)
        self.upsamplex4 = nn.Upsample(scale_factor=4, mode='bilinear')

        self.classifier = TwoLayerConv2d(in_channels=32, out_channels=output_nc)

        self.resnet_stages_num = resnet_stages_num

        self.if_upsample_2x = if_upsample_2x
        if self.resnet_stages_num == 5:
            layers = 512 * expand
        elif self.resnet_stages_num == 4:
            layers = 256 * expand
        elif self.resnet_stages_num == 3:
            layers = 128 * expand
        else:
            raise NotImplementedError
        self.conv_pred = nn.Conv2d(layers, 32, kernel_size=3, padding=1)

        self.output_sigmoid = output_sigmoid
        self.sigmoid = nn.Sigmoid()

    def forward(self, x1, x2):
        x1 = self.forward_single(x1)
        x2 = self.forward_single(x2)
        x = torch.abs(x1 - x2)
        if not self.if_upsample_2x:
            x = self.upsamplex2(x)
        x = self.upsamplex4(x)
        x = self.classifier(x)

        if self.output_sigmoid:
            x = self.sigmoid(x)
        return x

    def forward_single(self, x):
        # resnet layers
        x = self.resnet.conv1(x)
        x = self.resnet.bn1(x)
        x = self.resnet.relu(x)
        x = self.resnet.maxpool(x)

        x_4 = self.resnet.layer1(x) # 1/4, in=64, out=64
        x_8 = self.resnet.layer2(x_4) # 1/8, in=64, out=128

        if self.resnet_stages_num > 3:
            x_8 = self.resnet.layer3(x_8) # 1/8, in=128, out=256

        if self.resnet_stages_num == 5:
            x_8 = self.resnet.layer4(x_8) # 1/32, in=256, out=512
        elif self.resnet_stages_num > 5:
            raise NotImplementedError

        if self.if_upsample_2x:
            x = self.upsamplex2(x_8)
        else:
            x = x_8
        # output layers
        x = self.conv_pred(x)
        return x

class BASE_Transformer(ResNet):
    """
    Resnet of 8 downsampling + BIT + bitemporal feature Differencing + a small CNN
    """
    def __init__(self, input_nc, output_nc, with_pos, resnet_stages_num=5,
                 token_len=4, token_trans=True,
                 enc_depth=1, dec_depth=1,
                 dim_head=64, decoder_dim_head=64,
                 tokenizer=True, if_upsample_2x=True,
                 pool_mode='max', pool_size=2,
                 backbone='resnet18',
                 decoder_softmax=True, with_decoder_pos=None,
                 with_decoder=True):
        super(BASE_Transformer, self).__init__(input_nc, output_nc,backbone=backbone,
                                             resnet_stages_num=resnet_stages_num,
                                               if_upsample_2x=if_upsample_2x,
                                               )
        self.token_len = token_len
        self.conv_a = nn.Conv2d(32, self.token_len, kernel_size=1,
                                padding=0, bias=False)
        self.tokenizer = tokenizer
        if not self.tokenizer:
            #  if not use tokenzier，then downsample the feature map into a certain size
            self.pooling_size = pool_size
            self.pool_mode = pool_mode
            self.token_len = self.pooling_size * self.pooling_size

        self.token_trans = token_trans
        self.with_decoder = with_decoder
        dim = 32
        mlp_dim = 2*dim

        self.with_pos = with_pos
        if with_pos == 'learned':
            self.pos_embedding = nn.Parameter(torch.randn(1, self.token_len*2, 32))
        decoder_pos_size = 256//4
        self.with_decoder_pos = with_decoder_pos
        if self.with_decoder_pos == 'learned':
            self.pos_embedding_decoder =nn.Parameter(torch.randn(1, 32,
                                                                 decoder_pos_size,
                                                                 decoder_pos_size))
        self.enc_depth = enc_depth
        self.dec_depth = dec_depth
        self.dim_head = dim_head
        self.decoder_dim_head = decoder_dim_head
        self.transformer = Transformer(dim=dim, depth=self.enc_depth, heads=8,
                                       dim_head=self.dim_head,
                                       mlp_dim=mlp_dim, dropout=0)
        self.transformer_decoder = TransformerDecoder(dim=dim, depth=self.dec_depth,
                            heads=8, dim_head=self.decoder_dim_head, mlp_dim=mlp_dim, dropout=0,
                                                      softmax=decoder_softmax)

    def _forward_semantic_tokens(self, x):
        b, c, h, w = x.shape
        spatial_attention = self.conv_a(x)
        spatial_attention = spatial_attention.view([b, self.token_len, -1]).contiguous()
        spatial_attention = torch.softmax(spatial_attention, dim=-1)
        x = x.view([b, c, -1]).contiguous()
        tokens = torch.einsum('bln,bcn->blc', spatial_attention, x)
        return tokens

    def _forward_reshape_tokens(self, x):
        # b,c,h,w = x.shape
        if self.pool_mode == 'max':
            x = F.adaptive_max_pool2d(x, [self.pooling_size, self.pooling_size])
        elif self.pool_mode == 'ave':
            x = F.adaptive_avg_pool2d(x, [self.pooling_size, self.pooling_size])
        else:
            x = x
        tokens = rearrange(x, 'b c h w -> b (h w) c')
        return tokens

    def _forward_transformer(self, x):
        if self.with_pos:
            x += self.pos_embedding
        x = self.transformer(x)
        return x

    def _forward_transformer_decoder(self, x, m):
        b, c, h, w = x.shape
        if self.with_decoder_pos == 'fix':
            x = x + self.pos_embedding_decoder
        elif self.with_decoder_pos == 'learned':
            x = x + self.pos_embedding_decoder
        x = rearrange(x, 'b c h w -> b (h w) c')
        x = self.transformer_decoder(x, m)
        x = rearrange(x, 'b (h w) c -> b c h w', h=h)
        return x

    def _forward_simple_decoder(self, x, m):
        b, c, h, w = x.shape
        b, l, c = m.shape
        m = m.expand([h,w,b,l,c])
        m = rearrange(m, 'h w b l c -> l b c h w')
        m = m.sum(0)
        x = x + m
        return x

    def forward(self, x1, x2):
        # forward backbone resnet
        x1 = self.forward_single(x1)
        x2 = self.forward_single(x2)

        #  forward tokenzier
        if self.tokenizer:
            token1 = self._forward_semantic_tokens(x1)
            token2 = self._forward_semantic_tokens(x2)
        else:
            token1 = self._forward_reshape_tokens(x1)
            token2 = self._forward_reshape_tokens(x2)

        # forward transformer encoder
        if self.token_trans:
            self.tokens_ = torch.cat([token1, token2], dim=1)
            self.tokens = self._forward_transformer(self.tokens_)
            token1, token2 = self.tokens.chunk(2, dim=1)
        # forward transformer decoder
        if self.with_decoder:
            x1 = self._forward_transformer_decoder(x1, token1)
            x2 = self._forward_transformer_decoder(x2, token2)
        else:
            x1 = self._forward_simple_decoder(x1, token1)
            x2 = self._forward_simple_decoder(x2, token2)
        # feature differencing
        x = torch.abs(x1 - x2)
        if not self.if_upsample_2x:
            x = self.upsamplex2(x)
        x = self.upsamplex4(x)
        # forward small cnn
        x = self.classifier(x)
        if self.output_sigmoid:
            x = self.sigmoid(x)
        return x

class ConvReluBN(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3):
        super(ConvReluBN, self).__init__()
        self.layer = nn.Sequential(
            nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )
    def forward(self, x):
        return self.layer(x)

class ConvRelu(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3):
        super(ConvRelu, self).__init__()
        self.layer = nn.Sequential(
            nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, padding=1),
            nn.ReLU(inplace=True)
        )
    def forward(self, x):
        return self.layer(x)

class ChannelAttention(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3):
        super(ChannelAttention, self).__init__()

        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
        padding = 3 if kernel_size == 7 else 1
        self.conv1 = nn.Conv2d(in_channels*2, out_channels, kernel_size, padding=padding, bias=False)
        self.tanh = nn.Tanh()

    def forward(self, x_1, x_2):
        x = torch.cat([x_1, x_2], dim=1)
        x = self.conv1(x)
        return self.tanh(x)

class PreNorm(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.fn = fn
    def forward(self, x, **kwargs):
        return self.fn(self.norm(x), **kwargs)

class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout = 0.):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim),
            nn.Dropout(dropout)
        )
    def forward(self, x):
        return self.net(x)

class Attention(nn.Module):
    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
        super().__init__()
        inner_dim = dim_head *  heads
        self.heads = heads
        self.scale = dim ** -0.5

        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, dim),
            nn.Dropout(dropout)
        )

    def forward(self, x, mask = None):
        b, n, _, h = *x.shape, self.heads
        qkv = self.to_qkv(x).chunk(3, dim = -1)
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), qkv)
        dots = torch.einsum('bhid,bhjd->bhij', q, k) * self.scale
        mask_value = -torch.finfo(dots.dtype).max

        if mask is not None:
            mask = F.pad(mask.flatten(1), (1, 0), value = True)
            assert mask.shape[-1] == dots.shape[-1], 'mask has incorrect dimensions'
            mask = mask[:, None, :] * mask[:, :, None]
            dots.masked_fill_(~mask, mask_value)
            del mask

        attn = dots.softmax(dim=-1)
        out = torch.einsum('bhij,bhjd->bhid', attn, v)
        out = rearrange(out, 'b h n d -> b n (h d)')
        out = self.to_out(out)
        return out

class Residual(nn.Module):
    def __init__(self, fn):
        super().__init__()
        self.fn = fn
    def forward(self, x, **kwargs):
        return self.fn(x, **kwargs) + x

class Transformer(nn.Module):
    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout):
        super().__init__()
        self.layers = nn.ModuleList([])
        for _ in range(depth):
            self.layers.append(nn.ModuleList([
                Residual(PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout))),
                Residual(PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout)))
            ]))
    def forward(self, x, mask = None):
        for attn, ff in self.layers:
            x = attn(x, mask = mask)
            x = ff(x)
        return x

# Attention for bottleneck layers
class ChannelAttention_OnBottle(nn.Module):
    def __init__(self, in_planes, ratio=16, att_type='max'):
        super(ChannelAttention_OnBottle, self).__init__()
        self.att_type = att_type
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        # self.min_pool = nn.AdaptiveMinPool2d(1)

        self.fc1 = nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False)
        self.relu = nn.ReLU()
        self.fc2 = nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False)
        self.sigmoid = nn.Sigmoid()
        self.fc3 = nn.Linear(64, 512)
        self.fc4 = nn.Linear(96, 512)

    def forward(self, x):
        if self.att_type == 'max':
            out = self.fc2(self.relu(self.fc1(self.max_pool(x))))
        elif self.att_type == 'max_avg':
            max_out = self.relu(self.fc1(self.max_pool(x)))
            avg_out = self.relu(self.fc1(self.avg_pool(x)))
            out = torch.cat([max_out, avg_out], 1).squeeze()
            out = self.fc3(out).unsqueeze(0).unsqueeze(2).unsqueeze(3)
        elif self.att_type == 'avg_max_min':
            avg_out = self.relu(self.fc1(self.avg_pool(x)))
            # min_out = self.relu(self.fc1(self.min_pool(x)))
            max_out = self.relu(self.fc1(self.max_pool(x)))
            out = torch.cat([avg_out, min_out, max_out], 1).squeeze()
            out = self.fc4(out).unsqueeze(0).unsqueeze(2).unsqueeze(3)
        return self.relu(out)

class UNet_Change_Transformer(nn.Module):
    def __init__(self, pretrained=False, **kwargs):
        super(UNet_Change_Transformer, self).__init__()

        encoder_filters = [64, 64, 128, 256, 512]
        decoder_filters = np.asarray([48, 64, 96, 128, 320])

        self.encoder_filters = [64, 64, 128, 256, 512]
        self.decoder_filters = np.asarray([48, 64, 96, 160, 320])

        self.conv6 = ConvRelu(encoder_filters[-1], decoder_filters[-1])
        self.conv6_2 = ConvRelu(decoder_filters[-1] + encoder_filters[-2]*2 , decoder_filters[-1])
        self.conv7 = ConvRelu(decoder_filters[-1], decoder_filters[-2])
        self.conv7_2 = ConvRelu(decoder_filters[-2] + encoder_filters[-3]*2 , decoder_filters[-2])
        self.conv8 = ConvRelu(decoder_filters[-2], decoder_filters[-3])
        self.conv8_2 = ConvRelu(decoder_filters[-3] + encoder_filters[-4]*2 , decoder_filters[-3])
        self.conv9 = ConvRelu(decoder_filters[-3], decoder_filters[-4])
        self.conv9_2 = ConvRelu(decoder_filters[-4] + encoder_filters[-5]*2 , decoder_filters[-4])
        self.conv10 = ConvRelu(decoder_filters[-4], decoder_filters[-5])

        self.res = nn.Conv2d(decoder_filters[-5], 2, 1, stride=1, padding=0)

        self._initialize_weights()

        encoder = torchvision.models.resnet34(pretrained=pretrained)
        self.conv0 = ConvRelu(6,3)
        self.conv1 = nn.Sequential(
                        encoder.conv1,
                        encoder.bn1,
                        encoder.relu)
        self.conv2 = nn.Sequential(
                        encoder.maxpool,
                        encoder.layer1)
        self.conv3 = encoder.layer2
        self.conv4 = encoder.layer3
        self.conv5 = encoder.layer4

        self.ca_skip_5 = ChannelAttention(encoder_filters[-1], encoder_filters[-1])
        self.ca_skip_4 = ChannelAttention(encoder_filters[-2], encoder_filters[-2])
        self.ca_skip_3 = ChannelAttention(encoder_filters[-3], encoder_filters[-3])
        self.ca_skip_2 = ChannelAttention(encoder_filters[-4], encoder_filters[-4])
        self.ca_skip_1 = ChannelAttention(encoder_filters[-5], encoder_filters[-5])

        self.ca_bottle_max = ChannelAttention_OnBottle(512, att_type='max')
        self.ca_bottle_avg_min = ChannelAttention_OnBottle(512, att_type='max_avg')
        self.sigmoid = nn.Sigmoid()
        self.linearb = nn.Linear(1024, 512)

        dim = 64
        mlp_dim = 2*dim
        enc_depth = 2
        dim_head = 64
        decoder_dim_head = 64
        decoder_softmax = True
        self.transformer = Transformer(dim=dim, depth=3, heads=4,
                                       dim_head=dim_head,
                                       mlp_dim=mlp_dim, dropout=0.05)

        self.transformer_decoder = TransformerDecoder(dim=dim, depth=2,
                            heads=8, dim_head=decoder_dim_head, mlp_dim=mlp_dim, dropout=0,
                                                      softmax=decoder_softmax)
        self.classifier = TwoLayerConv2d(in_channels=512, out_channels=2)

    def forward_1(self, x1, x2):

        # Encoder 1
        x_1 = x1
        enc1_1 = self.conv1(x_1)
        enc2_1 = self.conv2(enc1_1)
        enc3_1 = self.conv3(enc2_1)
        enc4_1 = self.conv4(enc3_1)
        enc5_1 = self.conv5(enc4_1)

        # Encoder 2
        x_2 = x2
        enc1_2 = self.conv1(x_2)
        enc2_2 = self.conv2(enc1_2)
        enc3_2 = self.conv3(enc2_2)
        enc4_2 = self.conv4(enc3_2)
        enc5_2 = self.conv5(enc4_2)

        # Bottleneck
        enc5_1 = (self.ca_bottle_max(enc5_1)*enc5_1)
        enc5_2 = (self.ca_bottle_max(enc5_2)*enc5_2)

        enc5 = self.ca_skip_5(enc5_1,enc5_2)

        B_, C_, H_, W_ = enc5.shape
        enc5_i = enc5.view([B_, C_, H_*W_])
        enc5_i = self.transformer(enc5_i)
        enc5_i = enc5_i.view([B_, C_, H_, W_])
        enc5 = self.ca_skip_5(enc5_i,enc5)

        # Decoder
        enc4 = self.ca_skip_4(enc4_1, enc4_2)
        # enc4 = attention_block(enc4_1, enc4_2, self.encoder_filters[-2])
        dec6 = self.conv6(F.interpolate(enc5, scale_factor=2))
        dec6 = self.conv6_2(torch.cat([dec6, enc4
                ], 1))

        enc3 = self.ca_skip_3(enc3_1, enc3_2)
        # enc3 = attention_block(enc3_1, enc3_2, self.encoder_filters[-3])
        dec7 = self.conv7(F.interpolate(dec6, scale_factor=2))
        dec7 = self.conv7_2(torch.cat([dec7, enc3
                ], 1))

        enc2 = self.ca_skip_2(enc2_1, enc2_2)
        # enc2 = attention_block(enc2_1, enc2_2, self.encoder_filters[-4])
        dec8 = self.conv8(F.interpolate(dec7, scale_factor=2))
        dec8 = self.conv8_2(torch.cat([dec8, enc2
                ], 1))

        enc1 = self.ca_skip_2(enc1_1, enc1_2)
        # enc1 = attention_block(enc1_1, enc1_2, self.encoder_filters[-5])
        dec9 = self.conv9(F.interpolate(dec8, scale_factor=2))
        dec9 = self.conv9_2(torch.cat([dec9,
                enc1
                ], 1))

        dec10 = self.conv10(F.interpolate(dec9, scale_factor=2))
        out = self.res(dec10)
        return out

    def forward(self, x1, x2):
        # Encoder 1
        x_1 = x1
        enc1_1 = self.conv1(x_1)
        enc2_1 = self.conv2(enc1_1)
        enc3_1 = self.conv3(enc2_1)
        enc4_1 = self.conv4(enc3_1)
        enc5_1 = self.conv5(enc4_1)

        # Encoder 2
        x_2 = x2
        enc1_2 = self.conv1(x_2)
        enc2_2 = self.conv2(enc1_2)
        enc3_2 = self.conv3(enc2_2)
        enc4_2 = self.conv4(enc3_2)
        enc5_2 = self.conv5(enc4_2)

        # Bottleneck
        enc5_1 = (self.ca_bottle_max(enc5_1)*enc5_1)
        enc5_2 = (self.ca_bottle_max(enc5_2)*enc5_2)
        enc5_c = self.ca_skip_5(enc5_1,enc5_2)

        B_, C_, H_, W_ = enc5_c.shape
        enc5_i = enc5_c.view([B_, C_, H_*W_])
        enc5_i = self.transformer(enc5_i)
        enc5 = enc5_i.view([B_, C_, H_, W_])
        enc5 = self.ca_skip_5(enc5,enc5_c)

        # Decoder
        # enc4 = self.ca_skip_4(enc4_1, enc4_2)
        # enc4 = attention_block(enc4_1, enc4_2, self.encoder_filters[-2])
        dec6 = self.conv6(F.interpolate(enc5, scale_factor=2))
        dec6 = self.conv6_2(torch.cat([dec6, enc4_1, enc4_2
                ], 1))

        # enc3 = self.ca_skip_3(enc3_1, enc3_2)
        # enc3 = attention_block(enc3_1, enc3_2, self.encoder_filters[-3])
        dec7 = self.conv7(F.interpolate(dec6, scale_factor=2))
        dec7 = self.conv7_2(torch.cat([dec7, enc3_1, enc3_2
                ], 1))

        # enc2 = self.ca_skip_2(enc2_1, enc2_2)
        # enc2 = attention_block(enc2_1, enc2_2, self.encoder_filters[-4])
        dec8 = self.conv8(F.interpolate(dec7, scale_factor=2))
        dec8 = self.conv8_2(torch.cat([dec8, enc2_1, enc2_2
                ], 1))

        # enc1 = self.ca_skip_2(enc1_1, enc1_2)
        # enc1 = attention_block(enc1_1, enc1_2, self.encoder_filters[-5])
        dec9 = self.conv9(F.interpolate(dec8, scale_factor=2))
        dec9 = self.conv9_2(torch.cat([dec9, enc1_1, enc1_2], 1))

        dec10 = self.conv10(F.interpolate(dec9, scale_factor=2))

        out = self.res(dec10)

        # enc5_c = enc5_c.view([B_, C_, H_*W_])
        # interim_out = self.transformer_decoder(enc5_i, enc5_c)
        # interim_out = interim_out.view([B_, C_, H_, W_])
        # interim_out = self.classifier(enc5)

        return out #, interim_out

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d) or isinstance(m, nn.Linear):
                m.weight.data = nn.init.kaiming_normal_(m.weight.data)
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

class UNet_Change_Two_Transformer(nn.Module):
    def __init__(self, pretrained=False, **kwargs):
        super(UNet_Change_Two_Transformer, self).__init__()

        encoder_filters = [64, 64, 128, 256, 512]
        decoder_filters = np.asarray([48, 64, 96, 128, 320])

        self.encoder_filters = [64, 64, 128, 256, 512]
        self.decoder_filters = np.asarray([48, 64, 96, 160, 320])

        self.conv6 = ConvRelu(encoder_filters[-1], decoder_filters[-1])
        self.conv6_2 = ConvRelu(decoder_filters[-1] + encoder_filters[-2]*2, decoder_filters[-1])
        self.conv7 = ConvRelu(decoder_filters[-1], decoder_filters[-2])
        self.conv7_2 = ConvRelu(decoder_filters[-2] + encoder_filters[-3]*2, decoder_filters[-2])
        self.conv8 = ConvRelu(decoder_filters[-2], decoder_filters[-3])
        self.conv8_2 = ConvRelu(decoder_filters[-3] + encoder_filters[-4] , decoder_filters[-3])
        self.conv9 = ConvRelu(decoder_filters[-3], decoder_filters[-4])
        self.conv9_2 = ConvRelu(decoder_filters[-4] + encoder_filters[-5]*2 , decoder_filters[-4])
        self.conv10 = ConvRelu(decoder_filters[-4], decoder_filters[-5])

        self.res = nn.Conv2d(decoder_filters[-5], 2, 1, stride=1, padding=0)

        self._initialize_weights()

        encoder = torchvision.models.resnet34(pretrained=pretrained)
        self.conv1 = nn.Sequential(
                        encoder.conv1,
                        encoder.bn1,
                        encoder.relu)
        self.conv2 = nn.Sequential(
                        encoder.maxpool,
                        encoder.layer1)
        self.conv3 = encoder.layer2
        self.conv4 = encoder.layer3
        self.conv5 = encoder.layer4

        self.ca_skip_5 = ChannelAttention(encoder_filters[-1], encoder_filters[-1])
        # self.ca_skip_4 = ChannelAttention(encoder_filters[-2], encoder_filters[-2])
        # self.ca_skip_3 = ChannelAttention(encoder_filters[-3], encoder_filters[-3])
        # self.ca_skip_2 = ChannelAttention(encoder_filters[-4], encoder_filters[-4])
        self.ca_skip_1 = ChannelAttention(encoder_filters[-5], encoder_filters[-5])

        # self.ca_bottle_max = ChannelAttention_OnBottle(512, att_type='max')
        # self.ca_bottle_avg_min = ChannelAttention_OnBottle(512, att_type='max_avg')
        # self.sigmoid = nn.Sigmoid()
        # self.linearb = nn.Linear(1024, 512)

        dim = 64
        dim2 = 4096
        dim3 = 1024
        mlp_dim = 2*dim
        enc_depth = 2
        dim_head = 64
        decoder_dim_head = 64
        decoder_softmax = True
        self.transformer = Transformer(dim=dim, depth=3, heads=4,
                                       dim_head=dim_head,
                                       mlp_dim=mlp_dim, dropout=0.01)

        self.transformer2 = Transformer(dim=dim2, depth=2, heads=2,
                                       dim_head=dim_head,
                                       mlp_dim=dim2, dropout=0.001)

        # self.transformer3 = Transformer(dim=dim3, depth=2, heads=2,
        #                                dim_head=dim_head,
        #                                mlp_dim=dim3, dropout=0.001)

        self.classifier = TwoLayerConv2d(in_channels=512, out_channels=2)
        # self.convT = nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1)

    def forward(self, x1, x2):
        # Encoder 1
        x_1 = x1
        enc1_1 = self.conv1(x_1)
        enc2_1 = self.conv2(enc1_1)
        enc3_1 = self.conv3(enc2_1)
        enc4_1 = self.conv4(enc3_1)
        enc5_1 = self.conv5(enc4_1)

        # Encoder 2
        x_2 = x2
        enc1_2 = self.conv1(x_2)
        enc2_2 = self.conv2(enc1_2)
        enc3_2 = self.conv3(enc2_2)
        enc4_2 = self.conv4(enc3_2)
        enc5_2 = self.conv5(enc4_2)

        # Bottleneck
        # enc5_1 = (self.ca_bottle_max(enc5_1)*enc5_1)
        # enc5_2 = (self.ca_bottle_max(enc5_2)*enc5_2)
        # enc5_c = self.ca_skip_5(enc5_1,enc5_2)

        ## run 1: updating channel attention
        enc5 = self.ca_skip_5(enc5_1,enc5_2)
        B_, C_, H_, W_ = enc5.shape
        enc5_i = enc5.view([B_, C_, H_*W_]).contiguous()

        enc5_diff = (enc5_1 - enc5_2)
        spatial_attention = enc5_diff.view([B_, C_, H_*W_]).contiguous()
        spatial_attention = torch.softmax(spatial_attention, dim=-1)

        enc5_i = torch.einsum('bln,bln->bln', spatial_attention, enc5_i)
        enc5_t = self.transformer(enc5_i)
        enc5 = enc5_t.view([B_, C_, H_, W_]).contiguous()
        # enc5 = self.ca_skip_5(enc5_t,enc5)

        # Decoder
        # enc4 = self.ca_skip_4(enc4_1, enc4_2)

        dec6 = self.conv6(F.interpolate(enc5, scale_factor=2))
        dec6 = self.conv6_2(torch.cat([dec6, enc4_1, enc4_2], 1))

        # run4: depth=2, heads=2
        # enc3 = self.ca_skip_3(enc3_1, enc3_2)
        # B_, C_, H_, W_ = enc3.shape
        # enc3_i = enc3.view([B_, C_, H_*W_]).contiguous()

        # enc3_diff = (enc3_1 - enc3_2)
        # spatial_attention = enc3_diff.view([B_, C_, H_*W_]).contiguous()
        # spatial_attention = torch.softmax(spatial_attention, dim=-1)

        # enc3_i = torch.einsum('bln,bln->bln', spatial_attention, enc3_i)
        # enc3_t = self.transformer3(enc3_i)
        # enc3_t = enc3_t.view([B_, C_, H_, W_]).contiguous()
        # enc3 = self.ca_skip_3(enc3, enc3_t)

        dec7 = self.conv7(F.interpolate(dec6, scale_factor=2))
        dec7 = self.conv7_2(torch.cat([dec7, enc3_1, enc3_2], 1))

        ## run0: depth=2, heads=1
        enc2 = self.ca_skip_2(enc2_1, enc2_2)
        B_, C_, H_, W_ = enc2.shape
        enc2 = enc2.view([B_, C_, H_*W_]).contiguous()

        enc2_diff = (enc2_1 - enc2_2)
        spatial_attention = enc2_diff.view([B_, C_, H_*W_]).contiguous()
        spatial_attention = torch.softmax(spatial_attention, dim=-1)

        enc2 = torch.einsum('bln,bln->bln', spatial_attention, enc2)
        enc2 = self.transformer2(enc2)
        enc2 = enc2.view([B_, C_, H_,W_]).contiguous()

        dec8 = self.conv8(F.interpolate(dec7, scale_factor=2))
        dec8 = self.conv8_2(torch.cat([dec8, enc2], 1))

        # enc1 = self.ca_skip_2(enc1_1, enc1_2)
        dec9 = self.conv9(F.interpolate(dec8, scale_factor=2))
        dec9 = self.conv9_2(torch.cat([dec9, enc1_1, enc1_2], 1))

        dec10 = self.conv10(F.interpolate(dec9, scale_factor=2))
        out = self.res(dec10)

        return out #, interim_out

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d) or isinstance(m, nn.Linear):
                m.weight.data = nn.init.kaiming_normal_(m.weight.data)
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

class SiamUnet_conc(nn.Module):
    """SiamUnet_conc segmentation network."""

    def __init__(self, input_nbr, label_nbr):
        super(SiamUnet_conc, self).__init__()

        self.input_nbr = input_nbr

        self.conv11 = nn.Conv2d(input_nbr, 16, kernel_size=3, padding=1)
        self.bn11 = nn.BatchNorm2d(16)
        self.do11 = nn.Dropout2d(p=0.2)
        self.conv12 = nn.Conv2d(16, 16, kernel_size=3, padding=1)
        self.bn12 = nn.BatchNorm2d(16)
        self.do12 = nn.Dropout2d(p=0.2)

        self.conv21 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.bn21 = nn.BatchNorm2d(32)
        self.do21 = nn.Dropout2d(p=0.2)
        self.conv22 = nn.Conv2d(32, 32, kernel_size=3, padding=1)
        self.bn22 = nn.BatchNorm2d(32)
        self.do22 = nn.Dropout2d(p=0.2)

        self.conv31 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.bn31 = nn.BatchNorm2d(64)
        self.do31 = nn.Dropout2d(p=0.2)
        self.conv32 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
        self.bn32 = nn.BatchNorm2d(64)
        self.do32 = nn.Dropout2d(p=0.2)
        self.conv33 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
        self.bn33 = nn.BatchNorm2d(64)
        self.do33 = nn.Dropout2d(p=0.2)

        self.conv41 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.bn41 = nn.BatchNorm2d(128)
        self.do41 = nn.Dropout2d(p=0.2)
        self.conv42 = nn.Conv2d(128, 128, kernel_size=3, padding=1)
        self.bn42 = nn.BatchNorm2d(128)
        self.do42 = nn.Dropout2d(p=0.2)
        self.conv43 = nn.Conv2d(128, 128, kernel_size=3, padding=1)
        self.bn43 = nn.BatchNorm2d(128)
        self.do43 = nn.Dropout2d(p=0.2)

        self.upconv4 = nn.ConvTranspose2d(128, 128, kernel_size=3, padding=1, stride=2, output_padding=1)

        self.conv43d = nn.ConvTranspose2d(384, 128, kernel_size=3, padding=1)
        self.bn43d = nn.BatchNorm2d(128)
        self.do43d = nn.Dropout2d(p=0.2)
        self.conv42d = nn.ConvTranspose2d(128, 128, kernel_size=3, padding=1)
        self.bn42d = nn.BatchNorm2d(128)
        self.do42d = nn.Dropout2d(p=0.2)
        self.conv41d = nn.ConvTranspose2d(128, 64, kernel_size=3, padding=1)
        self.bn41d = nn.BatchNorm2d(64)
        self.do41d = nn.Dropout2d(p=0.2)

        self.upconv3 = nn.ConvTranspose2d(64, 64, kernel_size=3, padding=1, stride=2, output_padding=1)

        self.conv33d = nn.ConvTranspose2d(192, 64, kernel_size=3, padding=1)
        self.bn33d = nn.BatchNorm2d(64)
        self.do33d = nn.Dropout2d(p=0.2)
        self.conv32d = nn.ConvTranspose2d(64, 64, kernel_size=3, padding=1)
        self.bn32d = nn.BatchNorm2d(64)
        self.do32d = nn.Dropout2d(p=0.2)
        self.conv31d = nn.ConvTranspose2d(64, 32, kernel_size=3, padding=1)
        self.bn31d = nn.BatchNorm2d(32)
        self.do31d = nn.Dropout2d(p=0.2)

        self.upconv2 = nn.ConvTranspose2d(32, 32, kernel_size=3, padding=1, stride=2, output_padding=1)

        self.conv22d = nn.ConvTranspose2d(96, 32, kernel_size=3, padding=1)
        self.bn22d = nn.BatchNorm2d(32)
        self.do22d = nn.Dropout2d(p=0.2)
        self.conv21d = nn.ConvTranspose2d(32, 16, kernel_size=3, padding=1)
        self.bn21d = nn.BatchNorm2d(16)
        self.do21d = nn.Dropout2d(p=0.2)

        self.upconv1 = nn.ConvTranspose2d(16, 16, kernel_size=3, padding=1, stride=2, output_padding=1)

        self.conv12d = nn.ConvTranspose2d(48, 16, kernel_size=3, padding=1)
        self.bn12d = nn.BatchNorm2d(16)
        self.do12d = nn.Dropout2d(p=0.2)
        self.conv11d = nn.ConvTranspose2d(16, label_nbr, kernel_size=3, padding=1)

        self.sm = nn.LogSoftmax(dim=1)

    def forward(self, x1, x2):

        """Forward method."""
        # Stage 1
        x11 = self.do11(F.relu(self.bn11(self.conv11(x1))))
        x12_1 = self.do12(F.relu(self.bn12(self.conv12(x11))))
        x1p = F.max_pool2d(x12_1, kernel_size=2, stride=2)

        # Stage 2
        x21 = self.do21(F.relu(self.bn21(self.conv21(x1p))))
        x22_1 = self.do22(F.relu(self.bn22(self.conv22(x21))))
        x2p = F.max_pool2d(x22_1, kernel_size=2, stride=2)

        # Stage 3
        x31 = self.do31(F.relu(self.bn31(self.conv31(x2p))))
        x32 = self.do32(F.relu(self.bn32(self.conv32(x31))))
        x33_1 = self.do33(F.relu(self.bn33(self.conv33(x32))))
        x3p = F.max_pool2d(x33_1, kernel_size=2, stride=2)

        # Stage 4
        x41 = self.do41(F.relu(self.bn41(self.conv41(x3p))))
        x42 = self.do42(F.relu(self.bn42(self.conv42(x41))))
        x43_1 = self.do43(F.relu(self.bn43(self.conv43(x42))))
        x4p = F.max_pool2d(x43_1, kernel_size=2, stride=2)

        ####################################################
        # Stage 1
        x11 = self.do11(F.relu(self.bn11(self.conv11(x2))))
        x12_2 = self.do12(F.relu(self.bn12(self.conv12(x11))))
        x1p = F.max_pool2d(x12_2, kernel_size=2, stride=2)

        # Stage 2
        x21 = self.do21(F.relu(self.bn21(self.conv21(x1p))))
        x22_2 = self.do22(F.relu(self.bn22(self.conv22(x21))))
        x2p = F.max_pool2d(x22_2, kernel_size=2, stride=2)

        # Stage 3
        x31 = self.do31(F.relu(self.bn31(self.conv31(x2p))))
        x32 = self.do32(F.relu(self.bn32(self.conv32(x31))))
        x33_2 = self.do33(F.relu(self.bn33(self.conv33(x32))))
        x3p = F.max_pool2d(x33_2, kernel_size=2, stride=2)

        # Stage 4
        x41 = self.do41(F.relu(self.bn41(self.conv41(x3p))))
        x42 = self.do42(F.relu(self.bn42(self.conv42(x41))))
        x43_2 = self.do43(F.relu(self.bn43(self.conv43(x42))))
        x4p = F.max_pool2d(x43_2, kernel_size=2, stride=2)

        ####################################################
        # Stage 4d
        x4d = self.upconv4(x4p)
        pad4 = ReplicationPad2d((0, x43_1.size(3) - x4d.size(3), 0, x43_1.size(2) - x4d.size(2)))
        x4d = torch.cat((pad4(x4d), x43_1, x43_2), 1)
        x43d = self.do43d(F.relu(self.bn43d(self.conv43d(x4d))))
        x42d = self.do42d(F.relu(self.bn42d(self.conv42d(x43d))))
        x41d = self.do41d(F.relu(self.bn41d(self.conv41d(x42d))))

        # Stage 3d
        x3d = self.upconv3(x41d)
        pad3 = ReplicationPad2d((0, x33_1.size(3) - x3d.size(3), 0, x33_1.size(2) - x3d.size(2)))
        x3d = torch.cat((pad3(x3d), x33_1, x33_2), 1)
        x33d = self.do33d(F.relu(self.bn33d(self.conv33d(x3d))))
        x32d = self.do32d(F.relu(self.bn32d(self.conv32d(x33d))))
        x31d = self.do31d(F.relu(self.bn31d(self.conv31d(x32d))))

        # Stage 2d
        x2d = self.upconv2(x31d)
        pad2 = ReplicationPad2d((0, x22_1.size(3) - x2d.size(3), 0, x22_1.size(2) - x2d.size(2)))
        x2d = torch.cat((pad2(x2d), x22_1, x22_2), 1)
        x22d = self.do22d(F.relu(self.bn22d(self.conv22d(x2d))))
        x21d = self.do21d(F.relu(self.bn21d(self.conv21d(x22d))))

        # Stage 1d
        x1d = self.upconv1(x21d)
        pad1 = ReplicationPad2d((0, x12_1.size(3) - x1d.size(3), 0, x12_1.size(2) - x1d.size(2)))
        x1d = torch.cat((pad1(x1d), x12_1, x12_2), 1)
        x12d = self.do12d(F.relu(self.bn12d(self.conv12d(x1d))))
        x11d = self.conv11d(x12d)

        return self.sm(x11d)

class ResNet_UNet(torch.nn.Module):
    def __init__(self, input_nc, output_nc,
                 resnet_stages_num=5, backbone='resnet18',
                 output_sigmoid=False, if_upsample_2x=True):
        """
        In the constructor we instantiate two nn.Linear modules and assign them as
        member variables.
        """
        super(ResNet_UNet, self).__init__()
        expand = 1
        if backbone == 'resnet18':
            self.resnet = models.resnet18(pretrained=True, replace_stride_with_dilation=[False,True,True])
        elif backbone == 'resnet34':
            self.resnet = models.resnet34(pretrained=True, replace_stride_with_dilation=[False,True,True])
        else:
            raise NotImplementedError
        self.relu = nn.ReLU()
        self.upsamplex2 = nn.Upsample(scale_factor=2)
        self.upsamplex4 = nn.Upsample(scale_factor=4, mode='bilinear')

        self.resnet_stages_num = resnet_stages_num

        self.if_upsample_2x = if_upsample_2x
        if self.resnet_stages_num == 5:
            layers = 512 * expand
        elif self.resnet_stages_num == 4:
            layers = 256 * expand
        elif self.resnet_stages_num == 3:
            layers = 128 * expand
        else:
            raise NotImplementedError
        self.conv_pred = nn.Conv2d(384, 32, kernel_size=3, padding=1)

    def forward_single(self, x):
        # resnet layers
        x = self.resnet.conv1(x)
        x = self.resnet.bn1(x)
        x_2 = self.resnet.relu(x)
        x_2_pool = self.resnet.maxpool(x)

        x_4 = self.resnet.layer1(x_2_pool) # 1/4, in=64, out=64

        x_8 = self.resnet.layer2(x_4) # 1/8, in=64, out=128
        x_8_pool = self.resnet.maxpool(x_8)

        x_10 = self.resnet.layer3(x_8_pool) # 1/8, in=128, out=256

        if self.resnet_stages_num > 4:
            raise NotImplementedError

        # print(x_2.shape, x_4.shape, x_8.shape, x_10.shape)
        x = self.upsamplex2(x_10)

        return x_2, x_4, x_8, x_10

class BASE_Transformer_UNet(ResNet_UNet):
    """
    Resnet of 8 downsampling + BIT + bitemporal feature Differencing + a small CNN
    """
    def __init__(self, input_nc, output_nc, with_pos, resnet_stages_num=5,
                 token_len=4, token_trans=True,
                 enc_depth=1, dec_depth=1,
                 dim_head=64, decoder_dim_head=64,
                 tokenizer=True, if_upsample_2x=True,
                 pool_mode='max', pool_size=2,
                 backbone='resnet18',
                 decoder_softmax=True, with_decoder_pos=None,
                 with_decoder=True):
        super(BASE_Transformer_UNet, self).__init__(input_nc, output_nc,backbone=backbone,
                                             resnet_stages_num=resnet_stages_num,
                                               if_upsample_2x=if_upsample_2x,
                                               )

        print("using UNet Transformer !!!!")

        self.token_len = token_len
        self.tokenizer = tokenizer
        self.token_trans = token_trans
        self.with_decoder = with_decoder
        self.with_pos = with_pos

        if not self.tokenizer:
            #  if not use tokenzier，then downsample the feature map into a certain size
            self.pooling_size = pool_size
            self.pool_mode = pool_mode
            self.token_len = self.pooling_size * self.pooling_size

        # conv squeeze layers before transformer
        dim_5, dim_4, dim_3, dim_2 = 32, 32, 32, 32
        self.conv_squeeze_5 = nn.Sequential(nn.Conv2d(256, dim_5, kernel_size=1, padding=0, bias=False),
                                            nn.ReLU())
        self.conv_squeeze_4 = nn.Sequential(nn.Conv2d(128, dim_4, kernel_size=1, padding=0, bias=False),
                                            nn.ReLU())
        self.conv_squeeze_3 = nn.Sequential(nn.Conv2d(64, dim_3, kernel_size=1, padding=0, bias=False),
                                            nn.ReLU())
        self.conv_squeeze_2 = nn.Sequential(nn.Conv2d(64, dim_2, kernel_size=1, padding=0, bias=False),
                                            nn.ReLU())
        self.conv_squeeze_layers = [self.conv_squeeze_2, self.conv_squeeze_3, self.conv_squeeze_4, self.conv_squeeze_5]

        self.conv_token_5 = nn.Conv2d(dim_5, self.token_len, kernel_size=1, padding=0, bias=False)
        self.conv_token_4 = nn.Conv2d(dim_4, self.token_len, kernel_size=1, padding=0, bias=False)
        self.conv_token_3 = nn.Conv2d(dim_3, self.token_len, kernel_size=1, padding=0, bias=False)
        self.conv_token_2 = nn.Conv2d(dim_2, self.token_len, kernel_size=1, padding=0, bias=False)
        self.conv_tokens_layers = [self.conv_token_2, self.conv_token_3, self.conv_token_4, self.conv_token_5]

        self.conv_decode_5 = nn.Conv2d(64, 32, kernel_size=3, padding=1, bias=False)
        self.conv_decode_4 = nn.Conv2d(64, 32, kernel_size=3, padding=1, bias=False)
        self.conv_decode_3 = nn.Conv2d(64, 32, kernel_size=3, padding=1, bias=False)
        self.conv_decode_2 = nn.Conv2d(64, 32, kernel_size=3, padding=1, bias=False)
        self.conv_decode_layers = [self.conv_decode_2, self.conv_decode_3, self.conv_decode_4, self.conv_decode_5]

        if with_pos is 'learned':
            self.pos_embedding_5 = nn.Parameter(torch.randn(1, self.token_len*2, dim_5))
            self.pos_embedding_4 = nn.Parameter(torch.randn(1, self.token_len*2, dim_4))
            self.pos_embedding_3 = nn.Parameter(torch.randn(1, self.token_len*2, dim_3))
            self.pos_embedding_2 = nn.Parameter(torch.randn(1, self.token_len*2, dim_2))
            self.pos_embedding_layers = [self.pos_embedding_2, self.pos_embedding_3, self.pos_embedding_4, self.pos_embedding_5]

        decoder_pos_size = 256//4
        self.with_decoder_pos = with_decoder_pos
        if self.with_decoder_pos == 'learned':
            self.pos_embedding_decoder_5 =nn.Parameter(torch.randn(1, dim_5, 16, 16))
            self.pos_embedding_decoder_4 =nn.Parameter(torch.randn(1, dim_4, 32, 32))
            self.pos_embedding_decoder_3 =nn.Parameter(torch.randn(1, dim_3, 64, 64))
            self.pos_embedding_decoder_2 =nn.Parameter(torch.randn(1, dim_2, decoder_pos_size, decoder_pos_size))
            self.pos_embedding_decoder_layers = [self.pos_embedding_decoder_2, self.pos_embedding_decoder_3, self.pos_embedding_decoder_4, self.pos_embedding_decoder_5]

        self.enc_depth = enc_depth
        self.dec_depth = dec_depth
        self.dim_head = dim_head
        self.decoder_dim_head = decoder_dim_head
        self.transformer_5 = Transformer(dim=dim_5, depth=self.enc_depth, heads=4,
                                        dim_head=self.dim_head, mlp_dim=dim_5, dropout=0)
        self.transformer_decoder_5 = TransformerDecoder(dim=dim_5, depth=4, heads=4,
                                                        dim_head=self.decoder_dim_head, mlp_dim=dim_5, dropout=0, softmax=decoder_softmax)
        self.transformer_4 = Transformer(dim=dim_4, depth=self.enc_depth, heads=4,
                                        dim_head=self.dim_head, mlp_dim=dim_4, dropout=0)
        self.transformer_decoder_4 = TransformerDecoder(dim=dim_4, depth=4, heads=4, dim_head=self.decoder_dim_head,
                                                         mlp_dim=dim_4, dropout=0, softmax=decoder_softmax)
        self.transformer_3 = Transformer(dim=dim_3, depth=self.enc_depth, heads=8,
                                         dim_head=self.dim_head, mlp_dim=dim_3, dropout=0)
        self.transformer_decoder_3 = TransformerDecoder(dim=dim_3, depth=8, heads=8, dim_head=self.decoder_dim_head,
                                                        mlp_dim=dim_3, dropout=0, softmax=decoder_softmax)
        self.transformer_2 = Transformer(dim=dim_2, depth=self.enc_depth, heads=1,
                                         dim_head=32, mlp_dim=dim_2, dropout=0)
        self.transformer_decoder_2 = TransformerDecoder(dim=dim_2, depth=1, heads=1, dim_head=32,
                                                        mlp_dim=dim_2, dropout=0, softmax=decoder_softmax)
        self.transformer_layers = [self.transformer_2, self.transformer_3, self.transformer_4, self.transformer_5]
        self.transformer_decoder_layers = [self.transformer_decoder_2, self.transformer_decoder_3, self.transformer_decoder_4, self.transformer_decoder_5]

        self.conv_layer2_0 = TwoLayerConv2d(in_channels=128, out_channels=32, kernel_size=3)

        # # EXP NEW : Worked Better than BiT and changeformer
        self.conv_layer2 = nn.Sequential(nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, padding=1),
                                        nn.ReLU())
        self.conv_layer3 = nn.Sequential(nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, padding=1),
                                        nn.ReLU())
        self.conv_layer4 = nn.Sequential(nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, padding=1),
                                        nn.ReLU())
        self.classifier = nn.Conv2d(in_channels=32, out_channels=output_nc, kernel_size=3, padding=1)

        # FINALIZED CLASSIFIER: Works Best
        # self.classifier = nn.Conv2d(in_channels=32, out_channels=output_nc, kernel_size=3, padding=1)

        # # EXP NEW
        # self.conv_layer2 = nn.Sequential(nn.Conv2d(in_channels=40, out_channels=16, kernel_size=3, padding=1),
        #                                 nn.ReLU(),
        #                                 nn.Conv2d(in_channels=16, out_channels=8, kernel_size=3, padding=1),
        #                                 nn.ReLU())
        # self.conv_layer3 = nn.Sequential(nn.Conv2d(in_channels=48, out_channels=16, kernel_size=3, padding=1),
        #                                 nn.ReLU(),
        #                                 nn.Conv2d(in_channels=16, out_channels=8, kernel_size=3, padding=1),
        #                                 nn.ReLU())
        # self.conv_layer4 = nn.Sequential(nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, padding=1),
        #                                 nn.ReLU(),
        #                                 nn.Conv2d(in_channels=32, out_channels=16, kernel_size=3, padding=1),
        #                                 nn.ReLU())
        # self.classifier = nn.Sequential(nn.Conv2d(in_channels=8, out_channels=5, kernel_size=3, padding=1),
        #                                 nn.ReLU(),
        #                                 nn.Conv2d(in_channels=5, out_channels=output_nc, kernel_size=3, padding=1))

    def _forward_semantic_tokens(self, x, layer=None):
        b, c, h, w = x.shape
        spatial_attention = self.conv_tokens_layers[layer](x)
        spatial_attention = spatial_attention.view([b, self.token_len, -1]).contiguous()
        spatial_attention = torch.softmax(spatial_attention, dim=-1)
        x = x.view([b, c, -1]).contiguous()
        tokens = torch.einsum('bln,bcn->blc', spatial_attention, x)
        return tokens

    def _forward_transformer(self, x, layer):
        if self.with_pos:
            x += self.pos_embedding_layers[layer]
        x = self.transformer_layers[layer](x)
        return x

    def _forward_transformer_decoder(self, x, m, layer):
        b, c, h, w = x.shape
        if self.with_decoder_pos == 'learned':
            x = x + self.pos_embedding_decoder_layers[layer]
        x = rearrange(x, 'b c h w -> b (h w) c')
        x = self.transformer_decoder_layers[layer](x, m)
        x = rearrange(x, 'b (h w) c -> b c h w', h=h)
        return x

    def _forward_trans_module(self, x1, x2, layer):
        x1 = self.conv_squeeze_layers[layer](x1)
        x2 = self.conv_squeeze_layers[layer](x2)
        token1 = self._forward_semantic_tokens(x1, layer)
        token2 = self._forward_semantic_tokens(x2, layer)
        self.tokens_ = torch.cat([token1, token2], dim=1)
        self.tokens = self._forward_transformer(self.tokens_, layer)
        token1, token2 = self.tokens.chunk(2, dim=1)
        x1 = self._forward_transformer_decoder(x1, token1, layer)
        x2 = self._forward_transformer_decoder(x2, token2, layer)
        # return torch.abs(x1 - x2)

        # V1, V2
        # x1 = self._forward_transformer_decoder(x1, token2, layer)
        # x2 = self._forward_transformer_decoder(x2, token1, layer)
        # return torch.add(x1, x2)

        # V3
        diff_token = torch.abs(token2 - token1)
        diff_x = self.conv_decode_layers[layer](torch.cat([x1,x2], axis=1))
        x = self._forward_transformer_decoder(diff_x, diff_token, layer)
        return x

    def forward(self, x1, x2):
        # forward backbone resnet
        a_128, a_64, a_32, a_16 = self.forward_single(x1)
        b_128, b_64, b_32, b_16 = self.forward_single(x2)

        #  level 5 in=256x16x16 out=32x16x16
        x1, x2 = a_16, b_16
        out_5 = self._forward_trans_module(x1, x2, layer=3)
        out_5 = self.upsamplex2(out_5)

        # level 4: in=128x32x32 out=32x32x32
        x1, x2 = a_32, b_32
        out_4 = self._forward_trans_module(x1, x2, layer=2)
        out_4 = out_4 + out_5
        out_4 = self.upsamplex2(out_4)
        out_4 = self.conv_layer4(out_4)

        # level 3: in=64x64x64 out=32x64x64
        x1, x2 = a_64, b_64
        out_3 = self._forward_trans_module(x1, x2, layer=1)
        out_3 = out_3 + out_4
        # out_3 = self.conv_layer3(torch.cat([out_3, out_4], axis=1))
        out_3 = self.upsamplex2(out_3)
        out_3 = self.conv_layer3(out_3)

        # level 2: in=64x128x128
        out_2 = self.conv_layer2_0(torch.cat([a_128, b_128], 1))
        out_2 = out_2 + out_3
        # out_2 = self.conv_layer2(torch.cat([out_2, out_3], axis=1))
        out_2 = self.upsamplex2(out_2)
        out_2 = self.conv_layer2(out_2)

        # print(out_2.shape, out_3.shape, out_4.shape, out_5.shape)
        # forward small cnn
        x = self.classifier(out_2)

        return x

class Res34_Unet_Double(nn.Module):
    def __init__(self, pretrained=False, **kwargs):
        super(Res34_Unet_Double, self).__init__()

        encoder_filters = [64, 64, 128, 256, 512]
        decoder_filters = np.asarray([48, 64, 96, 160, 320])

        self.conv6 = ConvRelu(encoder_filters[-1], decoder_filters[-1])
        self.conv6_2 = ConvRelu(decoder_filters[-1] + encoder_filters[-2], decoder_filters[-1])
        self.conv7 = ConvRelu(decoder_filters[-1], decoder_filters[-2])
        self.conv7_2 = ConvRelu(decoder_filters[-2] + encoder_filters[-3], decoder_filters[-2])
        self.conv8 = ConvRelu(decoder_filters[-2], decoder_filters[-3])
        self.conv8_2 = ConvRelu(decoder_filters[-3] + encoder_filters[-4], decoder_filters[-3])
        self.conv9 = ConvRelu(decoder_filters[-3], decoder_filters[-4])
        self.conv9_2 = ConvRelu(decoder_filters[-4] + encoder_filters[-5], decoder_filters[-4])
        self.conv10 = ConvRelu(decoder_filters[-4], decoder_filters[-5])

        self.res = nn.Conv2d(decoder_filters[-5] * 2, 5, 1, stride=1, padding=0)
        #self.res = nn.Conv2d(decoder_filters[-5], 5, 1, stride=1, padding=0)
        self._initialize_weights()

        encoder = resnet34(pretrained=pretrained)
        self.conv0 = ConvRelu(6,3)
        self.conv1 = nn.Sequential(
                        encoder.conv1,
                        encoder.bn1,
                        encoder.relu)
        self.conv2 = nn.Sequential(
                        encoder.maxpool,
                        encoder.layer1)
        self.conv3 = encoder.layer2
        self.conv4 = encoder.layer3
        self.conv5 = encoder.layer4

    def forward1(self, x):
        batch_size, C, H, W = x.shape
        #x = self.conv0(x)
        enc1 = self.conv1(x)
        enc2 = self.conv2(enc1)
        enc3 = self.conv3(enc2)
        enc4 = self.conv4(enc3)
        enc5 = self.conv5(enc4)

        dec6 = self.conv6(F.interpolate(enc5, scale_factor=2))
        dec6 = self.conv6_2(torch.cat([dec6, enc4
                ], 1))

        dec7 = self.conv7(F.interpolate(dec6, scale_factor=2))
        dec7 = self.conv7_2(torch.cat([dec7, enc3
                ], 1))

        dec8 = self.conv8(F.interpolate(dec7, scale_factor=2))
        dec8 = self.conv8_2(torch.cat([dec8, enc2
                ], 1))

        dec9 = self.conv9(F.interpolate(dec8, scale_factor=2))
        dec9 = self.conv9_2(torch.cat([dec9,
                enc1
                ], 1))

        dec10 = self.conv10(F.interpolate(dec9, scale_factor=2))

        return dec10

    def forward(self, x1, x2):
        dec10_0 = self.forward1(x1)
        dec10_1 = self.forward1(x2)
        x = torch.cat([dec10_0, dec10_1], 1)
        #dec10 = self.sa(dec10) * dec10
        #x = self.forward1(x)
        return self.res(x)

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d) or isinstance(m, nn.Linear):
                m.weight.data = nn.init.kaiming_normal_(m.weight.data)
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

### Three layers UNet bottleneck -> siamese ###
''' # Encoder 1
        x_1 = x1
        enc1_1 = self.conv1(x_1)
        enc2_1 = self.conv2(enc1_1)
        enc3_1 = self.conv3(enc2_1)

        # Encoder 2
        x_2 = x2
        enc1_2 = self.conv1(x_2)
        enc2_2 = self.conv2(enc1_2)
        enc3_2 = self.conv3(enc2_2)

        enc3 = self.ca_skip_3(enc3_1, enc3_2)
        B_, C_, H_, W_ = enc3.shape
        enc3_i = enc3.view([B_, C_, H_*W_])
        enc3_i = self.transformer(enc3_i)
        enc3_i = enc3_i.view([B_, C_, H_, W_])
        enc3 = self.ca_skip_3(enc3_i,enc3)

        enc2 = self.ca_skip_2(enc2_1, enc2_2)
        dec8 = self.conv8(F.interpolate(enc3, scale_factor=2))
        dec8 = self.conv8_2(torch.cat([dec8, enc2
                ], 1))

        enc1 = self.ca_skip_2(enc1_1, enc1_2)
        dec9 = self.conv9(F.interpolate(dec8, scale_factor=2))
        dec9 = self.conv9_2(torch.cat([dec9,
                enc1
                ], 1))

        dec10 = self.conv10(F.interpolate(dec9, scale_factor=2))
        out = self.res(dec10)
        return out
'''


from random import randint

def gen_matrix():
    matrix = [
        #[1, 4, 4, 3, 1],
        #[3, 2, 5, 3, 4],
        #[2, 4, 2, 3, 5],
        #[5, 4, 3, 2, 1],
        #[3, 3, 4, 4, 3]
    ]
    for i in range(5):
        matrix.append([randint(1, 5) for i in range(5)])

    return matrix

def get_four_in_line(matrix):
    """ This methods loops around the matrix and finds if there are 4 numbers align verically/horizontally"""
    transposed_matrix = [[0 for i in range(5)] for j in range(5)]
    four_in_line_set = set() # we use a set so we dont have to put the 4-in-line items in order when testing

    for row in range(5):
        for col in range(5):
            transposed_matrix[row][col]=matrix[col][row]

    offset = 4
    for row in range(5):
        for col in range(2):
            #since matrix is 5x5 and we need to get 4-in-line. we only need to analize the first 2 items back and forth
            item = matrix[row][col]
            if matrix[row][col:offset + col] == [item, item + 1, item + 2, item + 3]: #checking 4 in line from 1st to 4th and 2nd to 5th element
                four_in_line_set.add(((row,col),(row,col + offset -1)))
            elif matrix[row][col:offset + col] == [item , item - 1, item -2, item-3]: #checking 4 in line from 1st to 4th and 2nd to 5th element backwards
                four_in_line_set.add(((row,col + offset -1 ),(row,col)))

    #we repeat the process with the transposed matrix
    for row in range(5):
        for col in range(2):
            item = transposed_matrix[row][col]
            if transposed_matrix[row][col:offset + col] == [item, item + 1, item + 2, item + 3]:
                four_in_line_set.add(((col,row),(col + offset -1,row)))
            elif transposed_matrix[row][col:offset + col] == [item , item - 1, item -2, item-3]:
                four_in_line_set.add(((col + offset -1, row),(col,row)))
    return four_in_line_set

def print_matrix(matrix):
    for row in matrix:
        print(row)

if __name__=='__main__':
    matrix = gen_matrix()
    print_matrix(matrix)

    four_in_line = get_four_in_line(matrix)

    if not four_in_line:
        print("there are no 4 consecutive numbers")
    else:
        print("lines found")
        for line in four_in_line:
            print(line)


from pydantic import BaseModel, Field,EmailStr

#method body
# # request_body = {
#     "metode_prediksi": "ML",
#     "range_prediksi":12,
#     "skala_prediksi":0,
#     "id":6,
#     "sampai": "2023-03-13T00:00"
# }

class method_prediksi(BaseModel):
    metode_prediksi: str=Field(default=None)
    range_prediksi: int= Field(default=None)
    skala_prediksi : int= Field(default=None)
    id : int= Field(default=None)
    sampai : str=Field(default=None)
    class Config:
        schema_extra={
            "post_demo":{
                "metode_prediksi":"menggunakan machine learning",
                "range_prediksi":"maksimal 12",
                " skala_prediksi":"skala 3",
                "id": "tergantung lokasi radar",
                "sampai": "2023-03-13T00:00"
            }
        }


import os, sys
import pickle
import numpy as np
import argparse

sys.path.append(".")
import data

def print_data(args, exp_name):

    # read outputs
    outputs = pickle.load(open(os.path.join(args.alg_output_root, exp_name, 'out.pk'), 'rb'))
    exp_args = outputs['args']
    results = outputs['results']

    # init
    key = exp_args.data.path
    if len(key) == 1:
        name = exp_args.model_ps.name[0]
    else:
        name = 'ACC'
    market_names = [k.split('/')[2].split('_')[0] for k in key]
    print('sources =', market_names)
    price_name = '/'.join(key[0].split('/')[1].split('_')[1:])
    #print(price_name)

    # process results
    K = results[0]['prediction_summary']['K']
    alpha_base = results[0]['prediction_summary']['alpha'][:K]
    alpha = np.sum(alpha_base)
    time = [r['time'] for r in results]
    # if len(key) == 1:
    #     price = [r['observation'][key[0]] for r in results]
    # else:
    #     print(results[0]['observation'])
    price = [[r['observation'][k] for k in key] for r in results]

    itv_min = [max(0, r['prediction_summary']['ps_updated'][0]) for r in results]
    itv_max = [min(max_val, r['prediction_summary']['ps_updated'][1]) for r in results]

    itv_bps = [r['prediction_summary']['base_ps']  for r in results]
    itv_bps = itv_bps[1:]

    # read TWAP results
    ds_TWAP = data.SinglePriceDataset(args.TWAP_data)
    price_TWAP = []
    for t in time:
        price_TWAP.append(1/ds_TWAP.read(t.astype(int)))
    price_TWAP = np.array(price_TWAP)

    args.num_data_max = min(args.num_data_max, len(time))
    if args.time_start:
        time_start = np.datetime64(args.time_start)
        for i, t in enumerate(time):
            if time_start <= t:
                data_start_index = i
                break
    else:
        data_start_index = 0
    time = time[data_start_index:args.num_data_max]
    price = price[data_start_index:args.num_data_max]
    price_TWAP = price_TWAP[data_start_index:args.num_data_max]

    itv_min = itv_min[data_start_index:args.num_data_max]
    itv_max = itv_max[data_start_index:args.num_data_max]
    itv_bps = itv_bps[data_start_index:args.num_data_max]

    if args.zoom_end_index == np.inf:
        args.zoom_end_index = len(time)
    time = time[args.zoom_start_index:args.zoom_end_index]
    price = price[args.zoom_start_index:args.zoom_end_index]
    price_TWAP = price_TWAP[args.zoom_start_index:args.zoom_end_index]

    itv_min = itv_min[args.zoom_start_index:args.zoom_end_index]
    itv_max = itv_max[args.zoom_start_index:args.zoom_end_index]
    itv_bps = itv_bps[args.zoom_start_index:args.zoom_end_index]

    # if args.show_base_ps:
    #     data_table = {}
    #     for time_i, price_i, itv_min_i, itv_max_i, itv_bps_i in zip(time, price, itv_min, itv_max, itv_bps):
    #         data_table[time_i] = {key[i]: p for i, p in enumerate(price_i)}
    #         data_table[time_i].update({'acc': [itv_min_i, itv_max_i]})

    #         out_str = f'[{time_i}] ACC = [{itv_min_i:.2f}, {itv_max_i:.2f}]'
    #         for k, v in data_table[time_i].items():
    #             if k is not 'acc':
    #                 key_name = k.split("/")[-1]
    #                 out_str += f', {key_name} = {v if v is not None else 0:.4f}'
    #                 out_str += f', BPS_{key_name} = [{itv_bps_i[k][0]:.2f}, {itv_bps_i[k][1]:.2f}]'
    #         print(out_str)

    # median
    price_obs = []
    for i, k in enumerate(key):
        price_i = [p[i] for p in price]
        price_obs.append(price_i)
    price_median = np.median(np.array(price_obs), 0)

    # time
    print('time &', ' & '.join([f'{t}' for t in time]))

    # price data
    for i, k in enumerate(key):
        price_i = [f'{p[i]:.4f}' for p in price]
        print(f'{k.split("/")[-1]} &', ' & '.join(price_i))

    # median
    print('median &', ' & '.join([f'{p:.4f}' for p in price_median]))

    # TWAP
    print('TWAP &', ' & '.join([f'{p:.4f}' for p in price_TWAP]))

    # BPS
    for i, k in enumerate(key):
        itvs = [f'[{itv[k][0]:.2f},{itv[k][1]:.2f}]' for itv in itv_bps]
        print(f'BPS_{k} =', ' & '.join(itvs))

    # ACC
    itv_ACC = []
    for mn, mx in zip(itv_min, itv_max):
        if mn == 0 and mx == max_val:
            itv_ACC.append('[$-\infty$,$\infty$]')
        else:
            itv_ACC.append(f'[{mn:.2f},{mx:.2f}]')
    print('ACon2 =', ' & '.join(itv_ACC))

if __name__ == '__main__':

    ## init a parser
    parser = argparse.ArgumentParser(description='online learning')
    parser.add_argument('--exp_names', type=str, nargs='+', default=
                        [
                            'three_sources_INV_ETH_SushiSwap_UniswapV2_coinbase_K_3_beta_1',
                            'two_sources_INV_ETH_SushiSwap_UniswapV2_K_2_beta_1',
                            'one_source_INV_ETH_SushiSwap_K_1_beta_0'
                        ]
    )
    parser.add_argument('--TWAP_data', type=str, default='data/price_ETH_INV/Keep3rV2SushiSwap')
    parser.add_argument('--output_root', type=str, default='output')
    parser.add_argument('--alg_output_root', type=str, default='output')
    parser.add_argument('--ours_name', type=str, default='ACon$^2$ (ours)')
    parser.add_argument('--fig_root', type=str, default='figs')
    parser.add_argument('--style', type=str, nargs='+', default=['-k', '-r', '-b'])
    parser.add_argument('--fontsize', type=int, default=15)
    parser.add_argument('--time_start', type=str, default='2022-04-02T11:03')
    parser.add_argument('--num_data_max', type=int, default=np.inf)
    parser.add_argument('--y_min', type=float, default=0)
    parser.add_argument('--y_max', type=float, default=10)
    parser.add_argument('--step', type=int, default=1)
    parser.add_argument('--log_scale', action='store_true')
    parser.add_argument('--tag', type=str)
    parser.add_argument('--zoom_start_index', type=int, default=4)
    parser.add_argument('--zoom_end_index', type=int, default=10)
    args = parser.parse_args()
    #color_list = ['C3', 'C4', 'C8', 'C9']
    color_list = ['r', 'b', 'gold', 'C9']
    max_val = 20

    for exp_name in args.exp_names:
        print(f'exp name = {exp_name}')
        print_data(args, exp_name)
        print()