import os
import tempfile
import math
import torch.distributed as dist
import numpy as np
import pandas as pd
import torch
import argparse
import torch.nn as nn
import torch.optim as optim
import torch.multiprocessing as mp
from apex.parallel import DistributedDataParallel as DDP
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import r2_score, mean_squared_error, mean_absolute_error
from LSTM_Model import STA_LSTM
from apex import amp
import random

torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '42594'
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3,4,5,6,7'
myseed = 42000

os.makedirs('../distributed_model', exist_ok=True)
def windows_input(input, output, window_range, pred_range):
    num = len(input) - window_range + 1 - pred_range
    input = input.unsqueeze(0)
    output = output.unsqueeze(0)
    seq_input = []
    for i in range(num):
        seq_input.append(input[:, i:i + window_range, :])
    input = torch.cat(seq_input, dim=0)
    seq_target = []
    for j in range(num):
        seq_target.append(output[:, window_range + j - 1:window_range + j - 1 + pred_range, :])
    output = torch.cat(seq_target, dim=0)
    return input, output

window_length=100
pred_length=10
data = pd.read_csv("../data/data_final.csv", index_col=0)
feature_num=len(data.columns)-1
data = data.values

input = torch.from_numpy(data[:,:feature_num]).float()
target = torch.from_numpy(data[:,feature_num:]).float()

config = {
    'epoch_num': 800,
    'batch_size': 50,
    'optim_hyper': {'lr': 0.0015, 'weight_decay': 0},
    'early_stop': 90,
    'optimizer': 'Adam',
    'save_name': 'sintering_model.pth'
}
loss_function = nn.MSELoss()
records = []

class dataset(Dataset):
    def __init__(self,
                 train_input,
                 train_target,
                 window_length,
                 pred_length,
                 mode='train'
                 ):
        self.mode = mode
        input_seq, output_seq = windows_input(train_input, train_target, window_length, pred_length)
        output_seq = output_seq.squeeze(-1)

        dev_split_start = round(len(input_seq) * 0.6)
        dev_split_end = dev_split_start + round(len(input_seq) * 0.2)
        if mode == 'train':
            indices = [i for i in range(dev_split_start) ]
        elif mode == 'dev':
            indices = [i for i in range(dev_split_start,dev_split_end) ]
        elif mode=='test':
            indices=[i for i in range(dev_split_end,len(input_seq)) ]
        self.data = input_seq[indices]
        self.target = output_seq[indices]

    def __getitem__(self, index):
        return self.data[index], self.target[index]
        # if self.mode in ['train', 'dev']:
        #     return self.data[index], self.target[index]
        # else:
        #     return self.data[index]

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

dataset_train = dataset(input,target,window_length,pred_length,mode='train')
dataset_dev=dataset(input,target,window_length,pred_length,mode='dev')
dataset_test=dataset(input,target,window_length,pred_length,mode='test')
test_dataloader = DataLoader(dataset=dataset_test,
                             batch_size=config['batch_size'],
                             shuffle=False,
                             drop_last=False,
                             pin_memory=True,
                             num_workers=0)

def set_up(rank, world_size):
    torch.cuda.set_device(rank)
    # 初始化某号进程组
    # 注：第一个参数选择后端，nccl后端是单机多卡情况下的推荐，比gloo快很多。
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def train_pro(train_dataloader, dev_dataloader, model, device, loss_fn, rank):
    early_point = 0
    min_loss = 10000
    # create model and move it to GPU with id rank
    model = model.cuda(rank)
    optimizer = getattr(torch.optim, config['optimizer'])(model.parameters(), **config['optim_hyper'])
    model, optimizer = amp.initialize(model, optimizer,
                                      opt_level='O1')
    ddp_model = DDP(model)
    for epoch in range(config['epoch_num']):
        ddp_model.train()
        for x, y in train_dataloader:
            optimizer.zero_grad()
            x, y = x.to(device), y.to(device)
            pred = ddp_model(x)
            loss = torch.sqrt(loss_fn(pred, y))
            with amp.scale_loss(loss, optimizer) as scaled_loss:
                scaled_loss.backward()
            optimizer.step()
        train_loss = dev_pro(dev_dataloader, ddp_model, device, loss_fn)
        records.append(train_loss)
        if (train_loss < min_loss):
            min_loss = train_loss
            if (rank == 0):
                print('Saving model (epoch = {:4d}, loss = {:.4f})'
                      .format(epoch + 1, min_loss))
            torch.save(ddp_model.state_dict(), '../distributed_model/{}'.format(config['save_name']))
            early_point = 0
        else:
            early_point += 1
        if early_point >= config['early_stop']:
            print("stop")
            if (rank == 0):
                loss_dev = pd.DataFrame(records)
                loss_dev.to_csv("loss_dev.csv")
            break
    # if rank == 0:
    #     print('Finished after {}'.format(epoch + 1))
    return min_loss

def dev_pro(dev_dataloader, model, device, loss_fn):
    model.eval()
    total_loss = 0
    for x, y in dev_dataloader:
        x, y = x.to(device), y.to(device)
        with torch.no_grad():
            pred = model(x)
            loss = loss_fn(pred, y)
        total_loss += loss.detach().cpu().item()
    total_loss = math.sqrt(total_loss / len(dev_dataloader))
    return total_loss

def test(tt_set, model, device):
    model.eval()  # set model to evalutation mode
    total_loss = torch.zeros(pred_length)
    preds = []
    for x, y in tt_set:  # iterate through the dataloader
        x = x.to(device)  # move data to device (cpu/cuda)
        y = y.to(device)
        with torch.no_grad():  # disable gradient calculation
            pred = model(x)  # forward pass (compute output)
            preds.append(pred.detach().cpu())  # collect prediction
            loss = torch.sum((pred - y) ** 2, dim=0)
        total_loss += loss.detach().cpu()
    total_loss = torch.sqrt(total_loss / len(tt_set.dataset))
    preds = torch.cat(preds, dim=0).numpy()  # concatenate all predictions and convert to a numpy array
    return preds, total_loss
def set_seed(seed):
    # 必须禁用模型初始化中的任何随机性。
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

def run_fn(rank, world_size):
    set_seed(myseed)
    print(f"Running on rank {rank}.")
    set_up(rank, world_size)

    train_sampler = DistributedSampler(dataset_train)
    dl_train_dist = DataLoader(dataset=dataset_train,
                               batch_size=config['batch_size'],
                               sampler=train_sampler,
                               shuffle=False)
    dev_sampler = DistributedSampler(dataset_dev)
    dl_dev_dist = DataLoader(dataset=dataset_dev,
                             batch_size=config['batch_size'],
                             sampler=dev_sampler,
                             shuffle=False)

    model = STA_LSTM(feature_num=feature_num, sa_hidden=window_length*2, ta_hidden=window_length,
                     output_size=pred_length, length=window_length)

    RMSE = train_pro(
        dl_train_dist,
        dl_dev_dist,
        model,
        device='cuda',
        loss_fn=loss_function,
        rank=rank
    )

def main(run_fn, world_size):
    mp.spawn(run_fn, args=(world_size,), nprocs=world_size, join=True)

if __name__ == "__main__":
    n_gpus = torch.cuda.device_count()
    assert n_gpus >= 2, f"Requires at least 2 GPUs to run, but got {n_gpus}"
    world_size = n_gpus

    set_seed(myseed)
    main(run_fn, world_size)
    model_1 = STA_LSTM(feature_num=feature_num, sa_hidden=40, ta_hidden=40,
                       output_size=pred_length, length=window_length).cuda()
    os.makedirs('../distributed_model', exist_ok=True)
    ckpt = torch.load('../distributed_model/{}'.format(config['save_name']), map_location='cuda')
    model_1.load_state_dict({k.replace('module.', ''): v for k, v in ckpt.items()})
    preds, loss = test(test_dataloader, model_1, "cuda")

    loss_test = pd.DataFrame(loss.numpy())
    preds=pd.DataFrame()

    # _, y_true = windows_input(da, data_test_target, lookback, time_range)
    # y_true = y_true.squeeze().numpy()
    # r2_loss = r2_score(y_true, preds, multioutput='raw_values')
    #
    # r2_loss_test = pd.DataFrame(r2_loss)
    # loss_test.to_csv("loss_test.csv")
    # r2_loss_test.to_csv("r2_loss_test.csv")


import inspect
import platform
from typing import Dict, Tuple, Union

import torch.nn as nn

from .registry import PLUGIN_LAYERS

if platform.system() == 'Windows':
    import regex as re  # type: ignore
else:
    import re  # type: ignore

def infer_abbr(class_type: type) -> str:
    """Infer abbreviation from the class name.

    This method will infer the abbreviation to map class types to
    abbreviations.

    Rule 1: If the class has the property "abbr", return the property.
    Rule 2: Otherwise, the abbreviation falls back to snake case of class
    name, e.g. the abbreviation of ``FancyBlock`` will be ``fancy_block``.

    Args:
        class_type (type): The norm layer type.

    Returns:
        str: The inferred abbreviation.
    """

    def camel2snack(word):
        """Convert camel case word into snack case.

        Modified from `inflection lib
        <https://inflection.readthedocs.io/en/latest/#inflection.underscore>`_.

        Example::

            >>> camel2snack("FancyBlock")
            'fancy_block'
        """

        word = re.sub(r'([A-Z]+)([A-Z][a-z])', r'\1_\2', word)
        word = re.sub(r'([a-z\d])([A-Z])', r'\1_\2', word)
        word = word.replace('-', '_')
        return word.lower()

    if not inspect.isclass(class_type):
        raise TypeError(
            f'class_type must be a type, but got {type(class_type)}')
    if hasattr(class_type, '_abbr_'):
        return class_type._abbr_  # type: ignore
    else:
        return camel2snack(class_type.__name__)

def build_plugin_layer(cfg: Dict,
                       postfix: Union[int, str] = '',
                       **kwargs) -> Tuple[str, nn.Module]:
    """Build plugin layer.

    Args:
        cfg (dict): cfg should contain:

            - type (str): identify plugin layer type.
            - layer args: args needed to instantiate a plugin layer.
        postfix (int, str): appended into norm abbreviation to
            create named layer. Default: ''.

    Returns:
        tuple[str, nn.Module]: The first one is the concatenation of
        abbreviation and postfix. The second is the created plugin layer.
    """
    if not isinstance(cfg, dict):
        raise TypeError('cfg must be a dict')
    if 'type' not in cfg:
        raise KeyError('the cfg dict must contain the key "type"')
    cfg_ = cfg.copy()

    layer_type = cfg_.pop('type')
    if layer_type not in PLUGIN_LAYERS:
        raise KeyError(f'Unrecognized plugin type {layer_type}')

    plugin_layer = PLUGIN_LAYERS.get(layer_type)
    abbr = infer_abbr(plugin_layer)

    assert isinstance(postfix, (int, str))
    name = abbr + str(postfix)

    layer = plugin_layer(**kwargs, **cfg_)

    return name, layer


# MatthewInkawhich
# Use this script to train a custom CNN without using LMDB for data layer input
# LINES TO MODIFY BEFORE RUNNING:
#   - Must manually set Configs

from __future__ import print_function
from matplotlib import pyplot
import numpy as np
import os
import shutil
import random
import skimage.io
import skimage.transform
import caffe2.python.predictor.predictor_exporter as pe
from caffe2.python import core, model_helper, net_drawer, workspace, visualize, brew, optimizer
from caffe2.proto import caffe2_pb2
from caffe2.python.predictor import mobile_exporter

########################################################################
# Configs
########################################################################
root_folder = os.path.join(os.path.expanduser('~'), 'DukeML', 'junk', 'mstar_NO_LMDB_output') #where bookkeeping files are outputted
save_trained_model_loc = root_folder
init_net_out = os.path.join(save_trained_model_loc, 'mstar_NO_LMDB_init_net.pb')
predict_net_out = os.path.join(save_trained_model_loc, 'mstar_NO_LMDB_predict_net.pb')
#training_lmdb = os.path.join(os.path.expanduser('~'), 'DukeML', 'datasets', 'mstar64', 'train_lmdb')
# validation_lmdb = os.path.join(os.path.expanduser('~'), 'DukeML', 'datasets', 'mstar64', 'validate_lmdb')
# testing_lmdb = os.path.join(os.path.expanduser('~'), 'DukeML', 'datasets', 'mstar64', 'test_lmdb')
TRAIN_DICTIONARY = os.path.join(os.path.expanduser('~'), 'DukeML', 'datasets', 'mstar64', 'training_labels')
VALIDATION_DICTIONARY = os.path.join(os.path.expanduser('~'), 'DukeML', 'datasets', 'mstar64', 'validation_labels')
TEST_DICTIONARY = os.path.join(os.path.expanduser('~'), 'DukeML', 'datasets', 'mstar64', 'testing_labels')
KEY_FILE = os.path.join(os.path.expanduser('~'), 'DukeML', 'datasets', 'mstar', 'MSTAR_KEYFILE')
num_classes = 8                   #number of image classes
training_net_batch_size = 30        #batch size for training
training_iters = 200               #training iterations
validation_images = 452            #total number of validation images
validation_interval = 25            #validate every ... training iterations
testing_images = 442               #total number of testing images
image_width = 64
image_height = 64
image_channels = 1

########################################################################
# create root_folder if not already there
if not os.path.isdir(root_folder):
    os.makedirs(root_folder)

# resetting workspace sets root_folder as "root" dir (bookkeeping files go here)
workspace.ResetWorkspace(root_folder)
print("workspace output folder:" + root_folder)

########################################################################
# Functions
########################################################################
def rescale(img, input_height, input_width):
    aspect = img.shape[1] / float(img.shape[0])
    if aspect > 1:
        return skimage.transform.resize(img, (input_width, int(aspect * input_height)))
    elif aspect < 1:
        return skimage.transform.resize(img, (int(input_width/aspect), input_height))
    else:
        return skimage.transform.resize(img, (input_width, input_height))

def crop_center(img, cropx, cropy):
    y, x, c = img.shape
    startx = x // 2 - (cropx // 2)
    starty = y // 2 - (cropy // 2)
    return img[starty:starty+cropy, startx:startx+cropx]

def prepare_image(img_path):
    img = skimage.io.imread(img_path)
    img = skimage.img_as_float(img)
    if (len(img.shape) < 3):
        img = np.expand_dims(img, axis=2)          # expand dims for greyscale images
    else:
        img = img[:, :, (2, 1, 0)]                 # RGB to BGR color order
    img = rescale(img, image_height, image_width)
    img = crop_center(img, image_height, image_width)
    img = img.swapaxes(1, 2).swapaxes(0, 1)    # HWC to CHW dimension
    #img = img * 255 - 128                      # Subtract mean = 128
    return img.astype(np.float32)

def make_batch(iterable, batch_size=1):
    length = len(iterable)
    for index in range(0, length, batch_size):
        yield iterable[index:min(index + batch_size, length)]

def extract_categories(key_file):
    key_dict = {}
    f = open(key_file)
    for line in f:
        l = line.split()
        key_dict[l[0]] = l[1]

    print(key_dict)
    return key_dict

class MSTAR_Dataset(object):
    def __init__(self, dictionary_file):
        self.categories = extract_categories(KEY_FILE)
        self.image_files = [line.split()[0] for line in open(dictionary_file)]
        self.labels = [line.split()[1] for line in open(dictionary_file)]

    def __getitem__(self, index):
        image = prepare_image(self.image_files[index])
        label = self.labels[index]
        return image, label

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

    def read(self, batch_size, shuffle=False):
        """Read (image, label) pairs in batch"""
        order = list(range(len(self)))
        if shuffle:
            random.shuffle(order)
        for batch in make_batch(order, batch_size):
            images, labels = [], []
            for index in batch:
                image, label = self[index]
                images.append(image)
                labels.append(label)
            yield np.stack(images).astype(np.float32), np.stack(labels).astype(np.int32).reshape((batch_size,))

def update_dims(height, width, kernel, stride, pad):
    new_height = ((height - kernel + 2*pad)//stride) + 1
    new_width = ((width - kernel + 2*pad)//stride) + 1
    return new_height, new_width

# define model architecture in terms of layers
def AddLeNetModel(model, data):
    # Image size: 64x64
    conv1 = brew.conv(model, data, 'conv1', dim_in=image_channels, dim_out=32, kernel=5)
    h,w = update_dims(height=image_height, width=image_width, kernel=5, stride=1, pad=0)
    # Image size: 60x60
    pool1 = brew.max_pool(model, conv1, 'pool1', kernel=2, stride=2)
    h,w = update_dims(height=h, width=w, kernel=2, stride=2, pad=0)
    relu1 = brew.relu(model, pool1, 'relu1')
    # Image size: 30x30
    conv2 = brew.conv(model, relu1, 'conv2', dim_in=32, dim_out=64, kernel=5)
    h,w = update_dims(height=h, width=w, kernel=5, stride=1, pad=0)
    # Image size: 26x26
    pool2 = brew.max_pool(model, conv2, 'pool2', kernel=2, stride=2)
    h,w = update_dims(height=h, width=w, kernel=2, stride=2, pad=0)
    relu2 = brew.relu(model, pool2, 'relu2')
    # Image size: 13x13
    conv3 = brew.conv(model, relu2, 'conv3', dim_in=64, dim_out=64, kernel=5)
    h,w = update_dims(height=h, width=w, kernel=5, stride=1, pad=0)
    # Image size: 9x9
    pool3 = brew.max_pool(model, conv3, 'pool3', kernel=2, stride=2)
    h,w = update_dims(height=h, width=w, kernel=2, stride=2, pad=0)
    relu3 = brew.relu(model, pool3, 'relu3')
    # Image size: 4x4
    # 50 * 4 * 4 stands for dim_out from previous layer multiplied by the image size
    ############# Change dim_in value if images are not 28x28
    fc3 = brew.fc(model, relu3, 'fc3', dim_in=64 * h * w, dim_out=500)
    fc3 = brew.relu(model, fc3, fc3)
    pred = brew.fc(model, fc3, 'pred', 500, num_classes)
    softmax = brew.softmax(model, pred, 'softmax')
    return softmax

# accuracy operator
def AddAccuracy(model, softmax, label):
    """Adds an accuracy op to the model"""
    accuracy = brew.accuracy(model, [softmax, label], "accuracy")
    return accuracy

# adds training operators to model
def AddTrainingOperators(model, softmax, label):
    """Adds training operators to the model."""
    xent = model.LabelCrossEntropy([softmax, label], 'xent')
    # compute the expected loss
    loss = model.AveragedLoss(xent, "loss")
    # track the accuracy of the model
    AddAccuracy(model, softmax, label)
    # use the average loss we just computed to add gradient operators to the model
    model.AddGradientOperators([loss])
    # do a simple stochastic gradient descent
    optimizer.build_sgd(
        model,
        base_learning_rate=0.1,
        policy="step",
        stepsize=1,
        gamma=0.999,
    )

# collects stats to inspect later
def AddBookkeepingOperators(model):
    """This adds a few bookkeeping operators that we can inspect later.

    These operators do not affect the training procedure: they only collect
    statistics and prints them to file or to logs.
    """
    # Print basically prints out the content of the blob. to_file=1 routes the
    # printed output to a file. The file is going to be stored under
    #     root_folder/[blob name]
    model.Print('accuracy', [], to_file=1)
    model.Print('loss', [], to_file=1)
    # Summarizes the parameters. Different from Print, Summarize gives some
    # statistics of the parameter, such as mean, std, min and max.
    for param in model.params:
        model.Summarize(param, [], to_file=1)
        model.Summarize(model.param_to_grad[param], [], to_file=1)
    # Now, if we really want to be verbose, we can summarize EVERY blob
    # that the model produces; it is probably not a good idea, because that
    # is going to take time - summarization do not come for free. For this
    # demo, we will only show how to summarize the parameters and their
    # gradients.

########################################################################
# Define training, testing, and deployment models
########################################################################
arg_scope = {"order": "NCHW"}
# Training model
train_model = model_helper.ModelHelper(
    name="train_net", arg_scope=arg_scope)
softmax = AddLeNetModel(train_model, "data")
AddTrainingOperators(train_model, softmax, "label")
#AddBookkeepingOperators(train_model)

# Validation model
val_model = model_helper.ModelHelper(
    name="val_net", arg_scope=arg_scope, init_params=False)
softmax = AddLeNetModel(val_model, "data")
AddAccuracy(val_model, softmax, "label")

# Testing model
test_model = model_helper.ModelHelper(
    name="test_net", arg_scope=arg_scope, init_params=False)
softmax = AddLeNetModel(test_model, "data")
AddAccuracy(test_model, softmax, "label")

# Deployment model
deploy_model = model_helper.ModelHelper(
    name="mstar_deploy", arg_scope=arg_scope, init_params=False)
AddLeNetModel(deploy_model, "data")
# You may wonder what happens with the param_init_net part of the deploy_model.
# No, we will not use them, since during deployment time we will not randomly
# initialize the parameters, but load the parameters from the db.

# Dump all protobufs to disk for later inspection
with open(os.path.join(root_folder, "train_net.pbtxt"), 'w') as fid:
    fid.write(str(train_model.net.Proto()))
with open(os.path.join(root_folder, "train_init_net.pbtxt"), 'w') as fid:
    fid.write(str(train_model.param_init_net.Proto()))
with open(os.path.join(root_folder, "test_net.pbtxt"), 'w') as fid:
    fid.write(str(test_model.net.Proto()))
with open(os.path.join(root_folder, "test_init_net.pbtxt"), 'w') as fid:
    fid.write(str(test_model.param_init_net.Proto()))
with open(os.path.join(root_folder, "deploy_net.pbtxt"), 'w') as fid:
    fid.write(str(deploy_model.net.Proto()))
print("Protocol buffers files have been created in your root folder: " + root_folder)

########################################################################
# Run training procedure
########################################################################
# Initialize
train_dataset = MSTAR_Dataset(TRAIN_DICTIONARY)
validation_dataset = MSTAR_Dataset(VALIDATION_DICTIONARY)
test_dataset = MSTAR_Dataset(TEST_DICTIONARY)

for image, label in train_dataset.read(batch_size=1):
    workspace.FeedBlob("data", image)
    workspace.FeedBlob("label", label)
    break

for image, label in validation_dataset.read(batch_size=1):
    workspace.FeedBlob("data", image)
    workspace.FeedBlob("label", label)
    break

for image, label in test_dataset.read(batch_size=1):
    workspace.FeedBlob("data", image)
    workspace.FeedBlob("label", label)
    break

# The parameter initialization network only needs to be run once.
workspace.RunNetOnce(train_model.param_init_net)
workspace.CreateNet(train_model.net, overwrite=True)
# initialize and create validation network
workspace.RunNetOnce(val_model.param_init_net)
workspace.CreateNet(val_model.net, overwrite=True)
# variables to track the accuracy & loss
accuracy = np.zeros(training_iters)
loss = np.zeros(training_iters)
# Now, we will manually run the network
for i in range(training_iters):
    for image, label in train_dataset.read(batch_size=training_net_batch_size, shuffle=True):
        workspace.FeedBlob("data", image)
        workspace.FeedBlob("label", label)
        #break
    workspace.RunNet(train_model.net)
    accuracy[i] = workspace.FetchBlob('accuracy')
    loss[i] = workspace.FetchBlob('loss')
    if (i % validation_interval == 0):
        print("Training iter: ", i)
        #run validation
        for image, label in validation_dataset.read(batch_size=validation_images, shuffle=True):
            workspace.FeedBlob("data", image)
            workspace.FeedBlob("label", label)
            #break
        workspace.RunNet(val_model.net.Proto().name)
        val_accuracy = workspace.FetchBlob('accuracy')
        print("Validation accuracy: ", str(val_accuracy))

# After the execution is done, let's plot the values.
pyplot.plot(loss, 'b')
pyplot.plot(accuracy, 'r')
pyplot.legend(('Loss', 'Accuracy'), loc='upper right')
pyplot.show()

########################################################################
# View data
########################################################################
pyplot.figure()
data = workspace.FetchBlob("data")
#_ = visualize.NCHW.ShowSingle(data[0][0])
_ = visualize.NCHW.ShowMultiple(data)
pyplot.figure()
softmax = workspace.FetchBlob('softmax')
_ = pyplot.plot(softmax[0], 'ro')
pyplot.title('Prediction for the first image')
pyplot.show()

########################################################################
# Run test on test net
########################################################################
#run a test pass on the test net
workspace.RunNetOnce(test_model.param_init_net)
workspace.CreateNet(test_model.net, overwrite=True)
for image, label in test_dataset.read(batch_size=testing_images, shuffle=True):
    workspace.FeedBlob("data", image)
    workspace.FeedBlob("label", label)
    #break
workspace.RunNet(test_model.net.Proto().name)
test_accuracy = workspace.FetchBlob('accuracy')
# After the execution is done, let's print the mean accuracy value.
print('\ntest_accuracy: %f' % test_accuracy)

########################################################################
# Save deploy model with trained weights and biases
########################################################################
# print("\n\nSaving the trained model to " + save_trained_model_loc)
# # Use the MOBILE EXPORTER to save the deploy model as pbs
# # https://github.com/caffe2/caffe2/blob/master/caffe2/python/predictor/mobile_exporter_test.py
workspace.RunNetOnce(deploy_model.param_init_net)
workspace.CreateNet(deploy_model.net, overwrite=True)
init_net, predict_net = mobile_exporter.Export(workspace, deploy_model.net, deploy_model.params)
with open(init_net_out, 'wb') as f:
    f.write(init_net.SerializeToString())
with open(predict_net_out, 'wb') as f:
    f.write(predict_net.SerializeToString())


import matplotlib.pyplot as plt
import os
import seaborn as sns
import numpy as np
import pingouin as pg

plt.rcParams["font.sans-serif"] = ["SimHei"]  # 用来正常显示中文标签

def model_overlap_plot(
    raw_data,
    sub_nums,
    assign_name,
    save_path=None,
    aspect_ratio=1.5,
    negative=False,
    print_corr=True,
    show_plot=False,
    title_fontsize=20,
    tick_fontsize=20,
):
    # Print the red text "===== Checking the consistency of the model capturing the experimental design ====="
    print("\033[31m===== 检测模型捕获实验设计的一致性 =====\033[0m")

    num_subs = len(sub_nums)  # Number of subjects
    num_models = len(assign_name)  # Number of models

    if save_path is not None:
        if not os.path.exists(save_path):
            os.makedirs(save_path)

    def value_data(value, negative=bool(negative)):
        return abs(1 - value) if negative else value

    # Calculate the new figure size based on the given aspect ratio
    new_fig_height = 5 * num_subs
    new_fig_width = new_fig_height * aspect_ratio

    # Create a new figure with subplots
    fig, axes = plt.subplots(
        num_subs, num_models, figsize=(new_fig_width, new_fig_height)
    )

    # Ensure axes is a 2D array even if there's only one row or one column
    if num_subs == 1:
        axes = axes.reshape(1, -1)
    if num_models == 1:
        axes = axes.reshape(-1, 1)

    # Loop through each subject and each model to populate the subplots
    for i, sub_num in enumerate(sub_nums):
        single_sub_data = raw_data[raw_data["sub_num"] == sub_num].reset_index(
            drop=True
        )

        # Extract and preprocess exp_design
        exp_design = single_sub_data["prop"].values
        if "MI" in exp_design:
            exp_design = [0.20 if x == "MI" else x for x in exp_design]
            exp_design = [0.80 if x == "MC" else x for x in exp_design]
        exp_design = [float(x) for x in exp_design]

        # Initialize a dictionary to store model estimates
        model_estimate = {}
        for names in assign_name:
            key = "_".join(names).split("_")[0:2]
            key = "_".join(key)
            model_estimate[key] = single_sub_data[names].mean(axis=1)

        # Plotting and correlation calculation
        colors = ["#e87e72", "#56bcc2"]
        for j, (key, value) in enumerate(model_estimate.items()):
            ax = axes[i, j]  # Select the corresponding subplot
            label = (
                "Bayesian Learning"
                if "bl" in key
                else "Reinforcement Learning"
                if "rl" in key
                else key
            )

            ax.plot(exp_design, color="darkblue", linewidth=5.0, label="exp_design")
            ax.plot(
                value_data(value, negative),
                color=colors[j],
                alpha=0.75,
                linewidth=5.0,
                label=label,
            )

            ax.set_title(
                f"Sub {sub_num}: {label}", fontsize=title_fontsize
            )  # Set y label fontsize
            ax.tick_params(axis="both", which="major", labelsize=tick_fontsize)
            ax.set_ylim([0, 1])

    plt.tight_layout()

    if print_corr:
        all_exp_design = raw_data["prop"].values
        if "MI" in all_exp_design:
            all_exp_design = [0.20 if x == "MI" else x for x in all_exp_design]
            all_exp_design = [0.80 if x == "MC" else x for x in all_exp_design]
        all_exp_design = [float(x) for x in all_exp_design]
        # Check for NaN values and raise an exception if found
        if np.isnan(exp_design).any() or np.isnan(value_data(value, negative)).any():
            raise ValueError(
                "Data contains NaN values, cannot calculate Pearson correlation."
            )

        for i in assign_name:
            # Calculate and print correlation and significance
            # 从 raw_data 中提取对应的所有 subject_num 的 value

            model_data = raw_data[i].mean(axis=1)
            value = np.array(model_data)

            if "bl" in i[0]:
                corr_label = "Bayesian Learning"
            elif "rl" in i[0]:
                corr_label = "Reinforcement Learning"

            # 计算 exp_design 和 values 的相关性
            # 使用 pingouin 进行统计分析
            # Convert the Series to a numpy array
            value_array = value_data(value, negative)

            # Calculate the correlation
            corr_results = pg.corr(all_exp_design, value_array)
            print(corr_results)

    # Save the figure to the specified path if given
    if save_path is not None:
        corr_results.to_csv(
            os.path.join(save_path, "correlation_results.csv"), index=False
        )
        plt.savefig(os.path.join(save_path, "model_overlap_plot.png"))
    if show_plot:
        plt.show()

    # Print the green text "========== Analysis completed =========="
    print("\033[32m========== 分析结束 ==========\033[0m")

# Update the function to include t-test for differences between volatility conditions
def plot_bl_v_boxplot(
    raw_data,
    subject_colname="sub_num",
    volatility_colname="volatile",
    bl_v_colname="bl_sr_v",
    colors=["#e87e72", "#56bcc2"],
    save_path=None,
    fig_size=(10, 6),
    print_ttest=True,
):
    print("\033[31m===== 贝叶斯模型箱线图绘制 =====\033[0m")
    # Convert volatility values to lowercase for consistency
    raw_data[volatility_colname] = raw_data[volatility_colname].str.lower()

    # Group the data by subject and volatility, then calculate the mean for each group
    grouped_data = (
        raw_data.groupby([subject_colname, volatility_colname])[bl_v_colname]
        .mean()
        .reset_index()
    )

    # Check if the required data for t-test is available
    if (
        grouped_data.empty
        or "s" not in grouped_data[volatility_colname].values
        or "v" not in grouped_data[volatility_colname].values
    ):
        raise ValueError("Insufficient data for t-test.")

    # Create the boxplot using seaborn
    plt.figure(figsize=(10, 6))
    sns.boxplot(x=volatility_colname, y=bl_v_colname, data=grouped_data, palette=colors)

    # Additional plot settings
    plt.title("Bayesian Learning")
    plt.xlabel("Volatility")
    plt.ylabel("Subject Mean Value")

    # Perform t-test between s and v conditions
    s_data = grouped_data[grouped_data[volatility_colname] == "s"][bl_v_colname]
    v_data = grouped_data[grouped_data[volatility_colname] == "v"][bl_v_colname]

    # 使用 pingouin 进行 t-test
    t_test_results = pg.ttest(s_data, v_data)

    if print_ttest:
        # 输出 APA 格式的报告
        print(t_test_results)

    if save_path is not None:
        # 保存 t-test 结果到 CSV 文件
        t_test_results.to_csv(
            os.path.join(save_path, "bl_ttest_results.csv"), index=False
        )
        # 保存 boxplot 图像
        fig = plt.gcf()
        fig.set_size_inches(fig_size[0], fig_size[1])
        plt.savefig(os.path.join(save_path, "bl_v_boxplot.png"), dpi=300)

    print("\033[32m========== 分析结束 ==========\033[0m")

# Re-plot the boxplot for RL model's alpha with all subjects included
def plot_rl_alpha_boxplot(
    alpha_table,
    subject_colname="sub_num",
    alpha_colnames=["alpha_s", "alpha_v"],
    colors=["#e87e72", "#56bcc2"],
    save_path=None,
    fig_size=(10, 6),
):
    print("\033[31m===== 强化学习模型箱线图绘制 =====\033[0m")
    # Check if the required columns exist in the dataframe
    for col in [subject_colname] + alpha_colnames:
        if col not in alpha_table.columns:
            print(f"Error: Missing required column '{col}' in the dataframe.")
            return

    # Check for missing values
    if alpha_table.isnull().any().any():
        print(
            "Warning: The dataframe contains missing values."
            "Please handle them before proceeding."
        )
        return

    # Reshape the data for boxplot
    melted_data = alpha_table.melt(
        id_vars=subject_colname,
        value_vars=alpha_colnames,
        var_name="volatile",
        value_name="alpha",
    )
    melted_data["volatile"] = melted_data["volatile"].map(
        {"alpha_s": "s", "alpha_v": "v"}
    )

    # Create the boxplot using seaborn
    plt.figure(figsize=(10, 6))
    sns.boxplot(x="volatile", y="alpha", data=melted_data, palette=colors)

    # Additional plot settings
    plt.title("Boxplot of RL Alpha")
    plt.xlabel("Volatility Condition")
    plt.ylabel("Alpha Value")

    # Perform t-test between s and v conditions
    s_data = alpha_table[next(col for col in alpha_colnames if "s" in col)]
    v_data = alpha_table[next(col for col in alpha_colnames if "v" in col)]
    # 使用 pingouin 进行 t-test
    t_test_results = pg.ttest(s_data, v_data)

    # 输出 APA 格式的报告
    print(t_test_results.round(3))

    if save_path is not None:
        t_test_results.to_csv(
            os.path.join(save_path, "rl_ttest_results.csv"), index=False
        )

        fig = plt.gcf()
        fig.set_size_inches(fig_size[0], fig_size[1])
        plt.savefig(os.path.join(save_path, "rl_alpha_boxplot.png"), dpi=300)

    print("\033[32m========== 分析结束 ==========\033[0m")


from django.conf.urls import url
from django.contrib.auth import views as auth_views

from . import views

urlpatterns = [
    url(r'^login$', auth_views.login, name='login'),
    url(r'^discourse_login', views.discourse_login, name='discourse_login'),
    url(r'^logout$', views.logout, name='logout'),
    url(r'^signup$', views.signup, name='signup'),

    url(r'^settings', views.settings),
    url(r'^config_dl', views.config_dl),
    url(r'^config', views.config),
    url(r'^logs', views.logs),
    url(r'^gift_code', views.gift_code),
    url(r'^trial', views.trial),
    url(r'^', views.index, name='index'),
]


import boto3
import logging
import hashlib
import os
import requests
from s3_client_lib.utils import *
from botocore.client import Config

logger = logging.getLogger(__name__)
import re

class S3Client:
    def __init__(
        self,
        address,
        access_key,
        secret_access_key,
        tenant=None,
        signature_version='s3v4',
    ):
        self.address = address
        config = Config(signature_version=signature_version)
        self.client = boto3.client(
            's3',
            endpoint_url=address,
            aws_access_key_id=access_key,
            aws_secret_access_key=secret_access_key,
            config=config,
        )

        def extend_bucket_with_tenant(
            request,
            signing_name,
            region_name,
            signature_version,
            request_signer,
            operation_name,
            **kwargs,
        ):
            try:
                logger.debug(f"BEFORE `{kwargs}`")
                logger.debug(f"BEFORE `{request.context}`")
                logger.info(f"BEFORE parsing bucket `{request.url}`")
                if tenant and request is not None:
                    bucket = None
                    match = re.match(r"(https|http):\/\/.+?\/(.*?)\/.*", request.url)
                    if match is not None:
                        bucket = match.group(2)
                    else:
                        # no bucket!!
                        logger.error(f"Cannot find bucket in url `{request}`, {kwargs}")
                        return
                    # bucket = request.context['signing'].get('bucket', kwargs.get('bucket', None))
                    logger.info(f' {tenant} => {bucket}')
                    request.url = request.url.replace(bucket, f'{tenant}%3A{bucket}')
                else:
                    logger.info('TENANT IS NOT DEFINED => do nothing')
            except Exception as e:
                logger.error(f'Error {e}')
        self.client.meta.events.register(
            'before-sign.s3.UploadPart', extend_bucket_with_tenant
        )
        self.resource = boto3.resource(
            's3',
            endpoint_url=address,
            aws_access_key_id=access_key,
            aws_secret_access_key=secret_access_key,
        )

    def finish_file_metadata(self, bucket, object_name, filename):
        s3_object = self.get_object_head(bucket, object_name)
        logger.debug(f'Fetched metadata header {s3_object}')
        self.update_metadata_object(
            bucket,
            object_name,
            {
                'contentLength': str(s3_object['ContentLength']),
                'contentType': s3_object['ContentType'],
                'contentName': filename,
            },
        )

    def create_bucket_if_not_exists(self, bucket_name, **kwargs):
        """
        Creates bucket in S3
        :param bucket_name:
        :param kwargs: {
            "ACL": 'private'|'public-read'|'public-read-write'|'authenticated-read',
            "Bucket": 'string',
            "CreateBucketConfiguration": {
                'LocationConstraint': 'EU'|'eu-west-1'|'us-west-1'|'us-west-2'|'ap-south-1'|'ap-southeast-1'|'ap-southeast-2'|'ap-northeast-1'|'sa-east-1'|'cn-north-1'|'eu-central-1'
            },
            "GrantFullControl"='string',
            "GrantRead"='string',
            "GrantReadACP"='string',
            "GrantWrite"='string',
            "GrantWriteACP"='string',
            "ObjectLockEnabledForBucket"=True|False
            }
        :return: dict
        """
        bucket = self.resource.Bucket(bucket_name)

        if bucket.creation_date:
            logger.info(f'The bucket exists {bucket_name}')
        else:
            logger.info(f'The bucket does not exist {bucket_name}')
            return self.client.create_bucket(Bucket=bucket_name, **kwargs)

    def upload_local_file(self, local_file, bucket, object_name):
        """

        :param local_file:
        :param object_name:
        :param size:
        :return:
        """
        # file is greater then 512 mb we have to upload it with multipart
        file_size = os.path.getsize(os.path.abspath(local_file))
        if file_size > GB_1:
            raise Exception('File size is greater than 1GB use S3MultipartClient')
        response_api = self.sign_s3_upload(bucket, object_name)
        s3_request_data = response_api['data']['fields']
        url = response_api['data']['url']
        filename = os.path.basename(local_file)
        with open(local_file, 'rb') as rf:
            logger.debug(f'Presigned upload => url: {url} fields: {s3_request_data}')
            s3_response = requests.post(
                url, data=s3_request_data, files={'file': (filename, rf)}
            )
            logger.debug(f'Response {s3_response} {s3_response.text}')
            self.finish_file_metadata(bucket, object_name, filename)
            return 'Upload completed'

    def sign_s3_upload(
        self, bucket, object_name, fields=None, conditions=None, expires=3600
    ) -> dict:
        """
        Create presigned url for upload of object to S3. This method is for smaller files.

        :param object_name: The name of the bucket to presign the post to
        :param fields:  A  dictionary  of  prefilled  form  fields  to  build  on  top  of.   Elements  that  may  be
                        included  are  acl,  Cache-Control,  Content-Type,  Content-Disposition, Content-Encoding,
                        Expires, success_action_redirect, redirect, success_action_status, and x-amz-meta-.
                        Note that if a particular element is included in the fields dictionary it will not
                        be automatically added to the conditions list.  You must specify a condition
                        for the element as well.
        :param conditions: A list of conditions to include in the policy.
                        Each element can be either a list or a structure.
                        For example:[{“acl”:  “public-read”}, [”content-length-range”, 2, 5],
                            [”starts-with”, “$success_action_redirect”, “”]]
                        Conditions that are included may pertain to acl,  content-length-range,  Cache-Control,
                        Content-Type,  Content-Disposition,  Content-Encoding,  Expires,  success_action_redirect,
                        redirect, success_action_status, and/or x-amz-meta-.Note that if you include a condition,
                        you must specify the a valid value in the fields dictionary as well.
                        A value will not be added automatically to the fields dictionary based on the conditions.
        :param expires: The number of seconds the presigned post is valid for.
        :return: {
            "data": {...},
            "url": string
        }
        """
        # object_name += "-${filename}"
        presigned_post = self.client.generate_presigned_post(
            Bucket=bucket,
            Key=object_name,
            Fields=fields,
            Conditions=conditions,
            ExpiresIn=expires,
        )
        url = os.path.join(os.path.join(self.address, bucket), object_name)
        return {'data': presigned_post, 'url': f'{url}'}

    def sign_s3_download(self, bucket, object_name, filename, expires=3600) -> dict:
        """
        Pre-signed url for download of file.
        :param bucket: The name of the bucket where is object
        :param object_name: The name of the object to presign
        :param filename: How should file named after download
        :param expires: The  number  of  seconds  the  presigned  url  is  valid  for.
                    By default it expires in an hour (3600 seconds)
        :return:{
            "url": string
        }
        """

        logger.debug(f'Signing url for {object_name} in bucket {bucket}')
        logger.info(f'Signing url for {object_name} in bucket {bucket}')

        presigned_url = self.client.generate_presigned_url(
            'get_object',
            Params={
                'Bucket': bucket,
                'Key': object_name,
                'ResponseContentDisposition': f'attachment; attachment; filename={filename}',
            },
            ExpiresIn=expires,
        )
        return {'url': presigned_url}

    def copy_from_s3(
        self, bucket, object_name, destination_path, chunk_size=CHUNK_SIZE_16M
    ):
        """
        Function will copy data from s3 to destination_path. For download is used get_object function which is dict which
        contains in Body key StreamingBody -> loading by chunks.

        :param object_name: The name of the object
        :param destination_path: destination where should object be downloaded
        :param chunk_size:
        :return:
        """

        logger.info(
            f'Copy data from S3 for: {object_name} from bucket: {bucket} to path: {destination_path}'
        )
        response = self.client.get_object(Bucket=bucket, Key=object_name)

        digest = hashlib.sha256()
        try:
            with open(destination_path, 'wb') as wf:
                for chunk in response['Body'].iter_chunks(chunk_size):
                    if chunk is None or not any(chunk):
                        break
                    digest.update(chunk)
                    wf.write(chunk)
            checksum_result = digest.hexdigest()
            logger.info(f'Checksum result: {checksum_result}')
            return checksum_result
        except Exception as e:
            logger.error(
                f'Something goes wrong in copy from S3 response: {response}, '
                f'destination = {destination_path}, '
                f'object_name: {object_name}'
                f'bucket: {bucket}'
                f'error: {e}'
            )

    def copy_to_s3(self, bucket, from_path, object_name, extra_args=None):
        """
        Copy file from local storage to S3.
        :param from_path: path to file
        :param object_name: how should object name
        :param extra_args: dict data which will be bind in object
        :return:
        """

        logger.info(
            f'Copy data to S3 for: {object_name} to bucket: {bucket} from path: {from_path}'
        )
        bucket = self.resource.Bucket(bucket)
        return bucket.upload_file(from_path, object_name, ExtraArgs=extra_args)

    def copy_from_bucket_to_bucket(
        self, object_name, source_bucket, destination_bucket
    ):
        """
        Copy objects from one bucket to another
        :param object_name: The name of the object to copy
        :param source_bucket:
        :param destination_bucket:
        :return:
        """
        copy_source = {'Bucket': source_bucket, 'Key': object_name}
        logger.info(
            f'Copy data between buckets S3 for: {object_name} source bucket: {source_bucket} '
            f'destination bucket: {destination_bucket}'
        )
        to_bucket = self.resource.Bucket(destination_bucket)
        return to_bucket.copy(copy_source, object_name)

    @staticmethod
    def __paginate(bucket, marker, size, iterator):
        result = []
        page = None

        while True:
            response_iterator = iterator(bucket, marker)

            for page in response_iterator:
                logger.debug(page)
                result.append(page)
            if page is None or page.get('Marker') == '':
                return page, result
            if size is not None:
                if len(result) > size:
                    return page.get('Marker', ''), result
            try:
                previous = marker
                marker = page.get('Marker', '')
            except KeyError:
                break
            if previous == marker and len(previous) > 0 and len(marker) > 0:
                break

    def list_objects(
        self, bucket=None, prefix=None, max_items=100, start_from=None, size=100
    ):
        """
        List objects from bucket

        :param prefix: filter objects which starts with
        :param max_items: max items on page
        :param start_from: marker from which will be data fetched
        :param size: max size of result data is is None return all data
        :return: marker, [objects]
        """
        config = {'PageSize': max_items}
        if prefix is not None:
            config.update({'Prefix': prefix})
        marker = start_from
        paginator = self.client.get_paginator('list_objects')

        def iterate(bucket, mark):
            config.update({'StartingToken': mark})
            return paginator.paginate(Bucket=bucket, PaginationConfig=config)

        return self.__paginate(bucket, marker, size, iterate)

    def search_objects(self, bucket, query, start_from=None, max_items=100, size=100):
        """
        Search by query in https://jmespath.org/ format.
        :param query:
        :param start_from: marker from start
        :param max_items: max items on page
        :param size: max size of result data is is None return all data
        :return: marker, [objects]
        """
        config = {'PageSize': max_items}
        marker = start_from
        paginator = self.client.get_paginator('list_objects')

        def iterate(bucket, mark):
            config.update({'StartingToken': mark})
            page_iterator = paginator.paginate(Bucket=bucket, PaginationConfig=config)
            return page_iterator.search(query)

        return self.__paginate(bucket, marker, size, iterate)

    def list_buckets(self):
        """
        List all buckets from S3
        :return:
        """
        return self.client.list_buckets()

    def delete_object(self, bucket, object_name):

        logger.debug(f'Deleting object {object_name} from bucket {bucket}')
        response = self.resource.Object(bucket, object_name).delete()
        return response

    def update_metadata_object(self, bucket, object_name, metadata):
        logger.debug(
            f'Update metadata on object {object_name} in bucket: {bucket} With metdata: {metadata}'
        )
        s3_object = self.resource.Object(bucket, object_name)
        s3_object.metadata.update(metadata)
        result = s3_object.copy_from(
            CopySource={'Bucket': bucket, 'Key': object_name},
            Metadata=s3_object.metadata,
            MetadataDirective='REPLACE',
        )
        return result

    def get_object_head(self, bucket, object_name):
        """
        Get mostly metadatada about object.
        :param object_name:
        :return:
        """

        response = self.client.head_object(Bucket=bucket, Key=object_name)
        return response

    def get_object(self, bucket, object_name):
        """
        Get whole object
        :param object_name:
        :return:
        """
        response = self.client.get_object(Bucket=bucket, Key=object_name)
        return response

    def get_stream(self, bucket_name, object_name):
        logger.info(
            f'Getting stream data from S3 for: {object_name} from bucket: {bucket_name}'
        )
        response = self.resource.Object(bucket_name=bucket_name, key=object_name)
        return S3File(response)

    def copy_data_from_s3_by_chunks_with_calc_sha(
        self,
        bucket,
        object_name,
        destination_path,
        chunk_size=CHUNK_SIZE_128M,
        process_chunk=lambda y, x: (y, x),
    ):
        """
        Function will copy data from s3 to destination_path. For download is used get_object function which is dict which
        contains in Body key StreamingBody -> loading by chunks.

        :param object_name:
        :param destination_path:
        :param bucket:
        :param chunk_size:
        :return:
        """
        import hashlib
        logger.info(
            f'Copy data from S3 for: {object_name} from bucket: {bucket} to path: {destination_path}'
        )
        response = self.s3_client.get_object(Bucket=bucket, Key=object_name)
        digest = hashlib.sha256()
        size = 0
        try:
            with open(destination_path, 'wb') as wf:
                for idx, chunk in enumerate(response['Body'].iter_chunks(chunk_size)):
                    if chunk is None or not any(chunk):
                        break
                    size += len(chunk)
                    digest.update(chunk)
                    process_chunk(idx, chunk)
                    wf.write(chunk)

            checksum_result = digest.hexdigest()
            logger.info(f'Checksum result: {checksum_result}')
            return checksum_result
        except Exception as e:
            print(e)
            logger.error(
                f'{self.name}-> Something goes wrong in copy from S3 response: {response}, '
                f'destination = {destination_path}, '
                f'object_name: {object_name}'
                f'bucket: {bucket}'
            )


from time import time
from typing import List, AnyStr
from collections import defaultdict

from dataset.COCO.coco_parser import COCOParser
from dataset.utils.random_buffer import RandomBuffer
from torch.utils.data import DataLoader, Dataset, RandomSampler

class COCOImages(COCOParser, Dataset):
    class ImgBuf:
        def __init__(self, img_ann, img_pix):
            self.img_ann, self.img_pix = img_ann, img_pix

    def _new_torch_img_buf(self, swap_img_idx):
        # swap_img_id = self.img_ids[swap_img_idx]
        # img_ann = self.instances.loadImgs(swap_img_id)[0]
        img_ann = self.img_anns[swap_img_idx]
        img_bin = self.load_img_binary(img_ann)
        torch_img = self.load_torch_pixels(img_bin)
        return COCOImages.ImgBuf(img_ann, torch_img)

    def _new_PIL_img_buf(self, swap_img_idx):
        # swap_img_id = self.img_ids[swap_img_idx]
        # img_ann = self.instances.loadImgs(swap_img_id)[0]
        img_ann = self.img_anns[swap_img_idx]
        img_bin = self.load_img_binary(img_ann)
        pil_img = self.load_PIL_pixels(img_bin)
        return COCOImages.ImgBuf(img_ann, pil_img)

    def __init__(self, split: AnyStr, cats=List[AnyStr], buffer_size=200, swap_prob_offset=0., show_stats=False, *args, **kwargs):
        super().__init__({split}, *args, **kwargs)

        if split == 'TRAIN':
            self.instances = self.instances_train2017
            self.captions = self.captions_train2017
        elif split == 'TEST':
            self.instances = self.image_info_test2017
            self.captions = self.image_info_test_dev2017
        elif split == 'VAL':
            self.instances = self.instances_val2017
            self.captions = self.captions_val2017

        self.img_ids = []
        cat_ids = self.instances.getCatIds(cats)
        self.label_map, self.rev_label_map = {}, {}
        for cat_id in cat_ids:
            self.img_ids.extend(self.instances.catToImgs[cat_id])
        self.img_anns = self.instances.loadImgs(self.img_ids)

        self.buffer = RandomBuffer(buffer_size, swap_prob_offset, show_stats)
        self.buffer.fill((self._new_PIL_img_buf(idx) for idx in range(COCOImages.__len__(self))))

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

    def __getitem__(self, idx):
        img_buf = self.buffer[lambda: self._new_PIL_img_buf(idx)]
        while img_buf.img_pix.mode != 'RGB':
            return COCOImages.__getitem__(self, idx+1)
        return img_buf.img_ann, img_buf.img_pix

    @staticmethod
    def collate_fn(elems):
        anns_batch, pixs_batch = [], []
        for ann, pixs in elems:
            anns_batch.append(ann)
            pixs_batch.append(pixs)
        return {'anns_batch': anns_batch, 'pixs_batch': pixs_batch}

if __name__ == '__main__':
    dataset = COCOImages('VAL', ['person', 'car'], 100, 0.0, show_stats=True)
    sampler = RandomSampler(dataset, True, 10000)
    loader = DataLoader(dataset=dataset,
                        batch_size=8,
                        sampler=sampler,
                        num_workers=1,
                        collate_fn=dataset.collate_fn,
                        pin_memory=False)
    t0 = time()
    print(len(dataset))
    for b, batch in enumerate(loader):
        print('Batch[{}]: {}'.format(loader.batch_size, b))
    print('b={}, s={}, b/s={}'.format(b, (time() - t0), b / (time() - t0)))


# import librarys

import gym
import random

env = gym.make('CartPole-v1')    # make an envirenment named 'Cartpole-v1' which has been defined in the library
goal_steps = 100

while True:
    obs = env.reset() # reset the status
    for i in range(goal_steps):
        obs, reward, done, info = env.step(random.randrange(0,2)) # set the random number which is the maximum numbers we can choose to make the bar uphold
        # obs: object current status, reward: reward according to the performance, done: bar is out of range or falled down
        if done:
            break # escape the while loop if done

        env.render() # render and display the environment


from telethon.errors import (
    BotMethodInvalidError,
    ChatSendInlineForbiddenError,
    ChatSendMediaForbiddenError,
)

from . import LOG_CHANNEL, LOGS, Button, asst, eor, get_string, ultroid_cmd

REPOMSG = """
â€¢ **ULTROID USERBOT** â€¢\n
â€¢ Repo - [Click Here](https://github.com/PredatorHackerzZ/Renamer-Bot)
â€¢ Addons - [Click Here](https://t.me/TellyFun_Official)
â€¢ Support - @TeleRoid14
"""

RP_BUTTONS = [
    [
        Button.url(get_string("bot_3"), "https://github.com/PredatorHackerzZ/Ultroid-UB"),
        Button.url("Addons", "https://t.me/TeleRoidGroup"),
    ],
    [Button.url("Support Group", "https://t.me/TeleRoid14")],
]

ULTSTRING = """ðŸŽ‡ **Thanks for Deploying Userbot!**

â€¢ Here, are the Some Basic stuff from, where you can Know, about its Usage."""

@ultroid_cmd(
    pattern="repo$",
    manager=True,
)
async def repify(e):
    try:
        q = await e.client.inline_query(asst.me.username, "")
        await q[0].click(e.chat_id)
        return await e.delete()
    except (
        ChatSendInlineForbiddenError,
        ChatSendMediaForbiddenError,
        BotMethodInvalidError,
    ):
        pass
    except Exception as er:
        LOGS.info("Error while repo command : " + str(er))
    await e.eor(REPOMSG)

@ultroid_cmd(pattern="ultroid$")
async def useUltroid(rs):
    button = Button.inline("Start >>", "initft_2")
    msg = await asst.send_message(
        LOG_CHANNEL,
        ULTSTRING,
        file="https://telegra.ph/file/cda94d1fd7e6519cdc1c9.jpg",
        buttons=button,
    )
    await eor(rs, f"**[Click Here]({msg.message_link})**")


#Decision tree
from sklearn.datasets import load_iris
from sklearn import tree
iris = load_iris()
print(iris.feature_names)
print(iris.data)
print(iris.target)
print(iris.target_names)
clf = tree.DecisionTreeClassifier()
clf = clf.fit(iris.data, iris.target)

import graphviz
dot_data = tree.export_graphviz(clf,out_file=None,
                         feature_names=iris.feature_names,
                         class_names=iris.target_names,
                         filled=True, rounded=True,
                         special_characters=True)
graph = graphviz.Source(dot_data)
graph.render("iris")
# http://scikit-learn.org/stable/modules/tree.html#decision-trees
# https://medium.com/towards-data-science/decision-trees-and-random-forests-for-classification-and-regression-pt-1-dbb65a458df

# Random Forrest
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import make_classification
X, y = make_classification(n_samples=1000, n_features=4,
                           n_informative=2, n_redundant=0,
                           random_state=0, shuffle=False)
clf = RandomForestClassifier(max_depth=2, random_state=0)
clf.fit(X, y)
# RandomForestClassifier(bootstrap=True, class_weight=None, criterion='gini',
#             max_depth=2, max_features='auto', max_leaf_nodes=None,
#             min_impurity_decrease=0.0, min_impurity_split=None,
#             min_samples_leaf=1, min_samples_split=2,
#             min_weight_fraction_leaf=0.0, n_estimators=10, n_jobs=1,
#             oob_score=False, random_state=0, verbose=0, warm_start=False)

# n_estimators=10 number of trees in forrest
#criterian  - quality of a split
# max_features - number of features to consider when looking for the best split:

print(clf.feature_importances_)
print(clf.predict([[0, 0, 0, 0]]))

# https://medium.com/machine-learning-101/chapter-5-random-forest-classifier-56dc7425c3e1
# http://www.kdnuggets.com/2016/09/reandom-forest-criminal-tutorial.html


#! /usr/bin/env python3
from typing import List
from random import randint

"""08/25/2019

Solution1:
I am very proud of this solution, because it is simple, short, and fast. Granted,
I did test run my solution against the OJ with some given test cases, but still
this is the very first complex DP solution I have written all by myself.

DP should be a very obvious choice for problems like this. However, I initially
wanted to analyze the string directly, which led to nowhere. A few obvious
observation:

1. the last substr must begin with the largest letter in the whole string
2. the last substr must continue to the end of the original string, because the
longer a string, the later it is ranked lexicologically.

This means, we will only need to record the position of the starting letter. The
DP idea is that if we have found the last substr of s[:i], how would such substr
change with an additional s[i+1]. Obviously, if s[i+1] is larger than the start
of the last substr of s[:i], we update the last substr of s[:i+1] to i+1. If
s[i+1] is the same as the start of the last substr of s[:i], and that s[i+1] is
not the same as s[i] (this is important, because a string of the same letters
can be treated as one letter represented by its first occurrence), we open up
a new possibility. This means we need an array to record multiple positions of
big letter, and our final result shall come from one of them.

The most complicated situation is when s[i+1] is smaller than the start of the
last substr of s[:i] or s[i+1] is part of a string of the same letter. If
currently we only have one candidate, we can ignore the new letter. But if we
already have more than one candidates, we need to compare s[i+1] to the
correspondingly positioned letter relative to all the candidates. For example,
suppose we have s[:i] = "zabza" and s[i+1] = 'c'. Our candidates = [0, 3]. Since
'c' is smaller than 'z', and we have two elements in candidates. We must compare
'c' to s[5 - 3 + 0] = s[2] = 'b'. Since 'c' > 'b', we know the last candidate
is bigger; we then update candidates to contain only the last elment. On the
other hand, if s[i+1] = 'a', since 'a' < 'b', we elimiate the last element. If
s[i+1] = 'b', we do nothing, because currently the two candidates lead the same
strings, and we need the next letter to decide.

The neat thing about this candidates array is that any position recorded in it
MUST lead the same exact string. This is because any difference would lead to
one candidate being kicked off. Therefore, when we do comparison with the newly
added letter, we only need to compare with letters led by candidates[0].

Another neat thing is that when the string is processed, our answer must come
from candidates[0] as it contains the longest string, even if there are other
candidates leading the same letters.

This solution has O(n) complexity, clocking in at 288ms, 74%
"""

class Solution:
    def lastSubstring(self, s: str) -> str:
        candidates: List[int] = [0]
        for i, le in enumerate(s[1:], start=1):
            if le > s[candidates[0]]:
                candidates = [i]
            # only start new if le is not following itself
            elif le == s[candidates[0]] and le != s[i - 1]:
                candidates.append(i)
            else:
                if len(candidates) > 1:
                    # new le bigger, keep the last candidate
                    if le > s[candidates[0] + i - candidates[-1]]:
                        candidates = candidates[-1:]
                    # new le smaller, pop the last candidate
                    elif le < s[candidates[0] + i - candidates[-1]]:
                        candidates.pop()
        return s[candidates[0] :]

sol = Solution()
s = "babb"
# print(sol.lastSubstring(s))

def gen_string(length: int):
    return "".join([chr(randint(97, 122)) for _ in range(length)])

print(gen_string(400000))


import argparse
import logging
import os
import sys

from tasks.valenco_translator_task import ValencoTranslatorTask

if __name__ == '__main__':
    logging.basicConfig(stream=sys.stdout, level=logging.INFO, format="%(message)s")
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', dest='input', required=False, help='Input file to process.')
    parser.add_argument('--output', dest='output', required=False, help='Output file to write results to.')
    parser.add_argument('--source_topic', dest='source_topic', required=False, help='source_topic')
    parser.add_argument('--sink_topic', dest='sink_topic', required=False, help='sink_topic')
    parser.add_argument('--kafka_brokers', dest='kafka_brokers', required=False, help='kafka_brokers')

    known_args, _ = parser.parse_known_args(sys.argv[1:])

    kafka_source_config = None
    if known_args.kafka_brokers and known_args.source_topic:
        kafka_source_config = {'bootstrap_servers': known_args.kafka_brokers, 'topic': known_args.source_topic}

    kafka_sink_config = None
    if known_args.kafka_brokers and known_args.sink_topic:
        kafka_sink_config = {'bootstrap_servers': known_args.kafka_brokers, 'topic': known_args.sink_topic}

    kafka_translator_task = ValencoTranslatorTask(
        jars_to_add=[f"file:///{os.path.dirname(os.path.realpath(__file__))}/connectors/flink-sql-connector-kafka-1.16.0.jar"],
        input_file=known_args.input,
        output_file=known_args.output,
        kafka_source_config=kafka_source_config,
        kafka_sink_config=kafka_sink_config
    )
    kafka_translator_task.translate()


from allennlp.data import DatasetReader, Token
from allennlp.data.dataset_readers.dataset_utils import span_utils
from allennlp.models.archival import load_archive
from allennlp.common.util import import_submodules
from bert_tokenization import BasicTokenizer
import ccg

class Predictor:
    def __init__(self, model_path):

        self.model_path = model_path
        self.tokenizer = BasicTokenizer(do_lower_case=False)

        archive = load_archive(model_path)
        self.model = archive.model
        config = archive.config.duplicate()
        dataset_reader_params = config["dataset_reader"]

        if "sentence_max_length" in dataset_reader_params.keys():
            del dataset_reader_params["sentence_max_length"]
        if "limit_files" in dataset_reader_params.keys():
            del dataset_reader_params["limit_files"]

        self.dataset_reader = DatasetReader.from_params(dataset_reader_params)
        self.model.eval()

    def predict_instance(self, original_text, ignore_clean=False):
        original_text = original_text.strip()
        print("original_txt:\n", original_text)
        text_tokens = self.tokenizer.tokenize(original_text, ignore_clean)
        print("text_tokens:\n", text_tokens)
        input_tokens = [Token(token) for token in text_tokens]

        instance = self.dataset_reader.text_to_instance(input_tokens)

        output = self.model.forward_on_instances([instance])[0]
        print("output[\"tags\"]:\n", output["tags"])
        spans = span_utils.bioul_tags_to_spans(output["tags"])
        print(spans)

        return format_text_annotation(original_text, text_tokens, spans)

def format_text_annotation(original_text, text_tokens, spans):
    assert len(original_text.strip()) == len(original_text)
    starts = []
    ends = []
    offset = 0

    for token in text_tokens:
        start_place = original_text.find(token, offset)
        assert start_place != -1

        starts.append(start_place)
        ends.append(start_place + len(token))
        assert original_text[starts[-1]: ends[-1]] == token
        offset = ends[-1]
    '''
        "tokenOffsets": [
            {
                "form": token,
                "startCharOffset": start,
                "endCharOffset": end
            } for (token, start, end) in zip(text_tokens, starts, ends)
        ],
    '''
    return {
        "corpusId": "",
        "id": "someFakeId",
        "text": original_text,
        "tokens": text_tokens,
        "sentences": {
            "generator": "UserSpecified",
            "score": 1.0,
            "sentenceEndPositions": [
                len(text_tokens)
            ]
        },
        "views": [
            {
                "viewName": "NER_CONLL",
                "viewData": [
                    {
                        "viewType": "edu.illinois.cs.cogcomp.core.datastructures.textannotation.View",
                        "viewName": "NER_CONLL",
                        "generator": "Ltf2TextAnnotation",
                        "score": 1.0,
                        "constituents": [
                            {
                                "label": tag,
                                "score": 1.0,
                                "start": start,
                                "end": end + 1
                            } for (tag, (start, end)) in spans
                        ]
                    },
                ]
            },
            {
                "viewName": "SENTENCE",
                "viewData": [
                    {
                        "viewType": "edu.illinois.cs.cogcomp.core.datastructures.textannotation.SpanLabelView",
                        "viewName": "SENTENCE",
                        "generator": "UserSpecified",
                        "score": 1.0,
                        "constituents": [
                            {
                                "label": "SENTENCE",
                                "score": 1.0,
                                "start": 0,
                                "end": len(text_tokens)
                            }
                        ]
                    }
                ]
            },
            {
                "viewName": "TOKENS",
                "viewData": [
                    {
                        "viewType": "edu.illinois.cs.cogcomp.core.datastructures.textannotation.TokenLabelView",
                        "viewName": "TOKENS",
                        "generator": "UserSpecified",
                        "score": 1.0,
                        "constituents": [
                            {
                                "label": "",
                                "score": 1.0,
                                "start": start,
                                "end": start + 1
                            } for start in range(len(text_tokens))
                        ]
                    }
                ]
            }
        ]
    }

def format_ner_output_to_json(original_text, text_tokens, ner_tags, spans):
    return {
        "text_tokens": text_tokens,
        "ner_tags": ner_tags,
        "spans": spans,
        "text_annotation": format_text_annotation(original_text, text_tokens, spans)
    }

def predict_single(text):
    predictor = Predictor("/scratch/hegler/lorelei-models/bert/som/exp1/model.tar.gz")
    return predictor.predict_instance(text)

def predict_with_initialized_predictor(predictor, text):
    return predictor.predict_instance(text)

if __name__ == '__main__':
    text = "DHEGAYSO+Sawiro: Qaraxii Muqdisho Ruxay , wafdigii ku sugnaa Iyo qorshaha Imaatinka madaxwaynaha Turkey ( VOA & Reuters )"
    print(predict_single(text))
    # if somehow can keep a python instance running, then:
    # load the predictor first, as it takes time
    predictor = Predictor("/scratch/hegler/lorelei-models/bert/som/exp1/model.tar.gz")
    for query in range(5):
        print(predict_with_initialized_predictor(predictor, text))


import random

from datetime import datetime, timedelta
from django.db import transaction
from django.db.models import OuterRef, Exists

from common.models import *
from config import *

class TimeUtil(object):
    @classmethod
    def format_time(cls, date_string):
        return datetime.strptime(date_string, "%Y-%m-%d %H:%M:%S")

class MoneyUtil(object):
    @classmethod
    def yuan_to_fen(cls, yuan):
        return round(100*yuan)

class PaymentUtil(object):
    @classmethod
    def gen_serial_number(cls, t_time):
        return t_time.strftime('%Y%m%d%H%M%S%f') + ('%012d' % random.randint(0, 999999999999))

class MockProcessUtil(object):

    @classmethod
    def mock_get_coupon(cls, client, coupon_rule, obtain_time, originator_merchant):
        new_coupon = Coupon.objects.create(
            rule=coupon_rule,
            client=client,
            discount=coupon_rule.discount,
            min_charge=coupon_rule.min_charge,
            originator_merchant=originator_merchant,
            status=COUPON_STATUS.NOT_USED,
            obtain_datetime=TimeUtil.format_time(obtain_time)
        )

        return new_coupon

    @classmethod
    def mock_request_payment(cls, user_id=None, merchant_id=None, coupon_id=None, order_price=0,
                            pay_time_str=None, pay_way=None):

        client = Client.objects.get(pk=user_id)
        merchant = Merchant.objects.get(pk=merchant_id)
        pay_time = TimeUtil.format_time(pay_time_str)

        coupon = None
        if coupon_id:
            coupon = Coupon.objects.get(pk=coupon_id)
            coupon.status = COUPON_STATUS['USED']
            coupon.use_datetime = pay_time
            coupon.save()

        paid_price = order_price - coupon.discount if coupon else order_price

        new_payment = Payment.objects.create(
            serial_number=PaymentUtil.gen_serial_number(pay_time),
            pay_channel=pay_way, # PAY_CHANNELS['WECHAT'],
            status=PAYMENT_STATUS['UNPAID'],
            merchant=merchant,
            client=client,
            coupon=coupon,
            order_price=order_price,
            datetime=pay_time,
            platform_share=round(max(0, paid_price) * PLATFORM_SHARE_RATE) if coupon else 0,
            inviter_share=round(max(0, paid_price) * INVITER_SHARE_RATE) if coupon else 0,
            originator_share=round(max(0, paid_price) * ORIGINATOR_SHARE_RATE) if coupon else 0,
        )

        return new_payment

    @classmethod
    def mock_on_payment_success(cls, platform_account_id, payment, time_str):
        """ mock_payment """
        with transaction.atomic():
            platform_account = Account.objects.select_for_update().get(id=platform_account_id)
            merchant_account = Account.objects.select_for_update().get(
                id=payment.merchant.account.id)
            payment = Payment.objects.select_for_update().get(serial_number=payment.serial_number)
            timestamp = TimeUtil.format_time(time_str)

            if payment.status != PAYMENT_STATUS['UNPAID']:
                raise Exception("payment's status is invalid") # Duplicated callback
            payment.status = PAYMENT_STATUS['FROZEN']
            payment.save()

            is_wechat_pay = True if payment.pay_channel == PAY_CHANNELS['WECHAT'] else False

            if not payment.coupon:  # 没有使用优惠券的情况
                merchant_account.balance = merchant_account.balance + payment.order_price
                merchant_account.save()

                Transaction.objects.bulk_create([
                    Transaction(
                        content_object=payment,
                        transaction_type=TRANSACTION_TYPE['PLATFORM_RECEIVE'],
                        datetime=timestamp,
                        account=platform_account,
                        amount=payment.order_price,
                        balance_after_transaction=(platform_account.balance if is_wechat_pay
                                                   else platform_account.alipay_balance) + payment.order_price
                    ),
                    Transaction(
                        content_object=payment,
                        transaction_type=TRANSACTION_TYPE['PLATFORM_EXPEND_MERCHANT_RECEIVE'],
                        datetime=timestamp,
                        account=platform_account,
                        amount=-payment.order_price,
                        balance_after_transaction=platform_account.balance if is_wechat_pay else platform_account.alipay_balance
                    ),
                    Transaction(
                        content_object=payment,
                        transaction_type=TRANSACTION_TYPE['MERCHANT_RECEIVE'],
                        datetime=timestamp,
                        account=merchant_account,
                        amount=payment.order_price,
                        balance_after_transaction=merchant_account.balance if is_wechat_pay else merchant_account.alipay_balance
                    )
                ])
            else:  # 使用了优惠券的情况
                paid_price = payment.order_price - payment.coupon.discount
                total_share = payment.platform_share + payment.inviter_share + payment.originator_share

                if is_wechat_pay:
                    platform_account.balance = platform_account.balance + total_share
                    merchant_account.balance = merchant_account.balance + paid_price - total_share
                else:
                    platform_account.alipay_balance = platform_account.alipay_balance + total_share
                    merchant_account.alipay_balance = merchant_account.alipay_balance + paid_price - total_share
                platform_account.save()
                merchant_account.save()

                Transaction.objects.bulk_create([
                    Transaction(
                        content_object=payment,
                        transaction_type=TRANSACTION_TYPE['PLATFORM_RECEIVE'],
                        datetime=timestamp,
                        account=platform_account,
                        amount=paid_price,
                        balance_after_transaction=(platform_account.balance if is_wechat_pay
                                                   else platform_account.alipay_balance) + paid_price - total_share
                    ),
                    Transaction(
                        content_object=payment,
                        transaction_type=TRANSACTION_TYPE['PLATFORM_EXPEND_MERCHANT_RECEIVE'],
                        datetime=timestamp,
                        account=platform_account,
                        amount=-(paid_price - total_share),
                        balance_after_transaction=platform_account.balance if is_wechat_pay else platform_account.alipay_balance
                    ),
                    Transaction(
                        content_object=payment,
                        transaction_type=TRANSACTION_TYPE['MERCHANT_RECEIVE'],
                        datetime=timestamp,
                        account=merchant_account,
                        amount=paid_price - total_share,
                        balance_after_transaction=merchant_account.balance if is_wechat_pay else merchant_account.alipay_balance
                    )
                ])

    @classmethod
    def mock_request_refund(cls, payment, time_str):
        with transaction.atomic():
            payment = Payment.objects.select_for_update().get(serial_number=payment.serial_number)

            if payment.status == PAYMENT_STATUS['FROZEN']:
                new_refund = Refund.objects.create(
                    serial_number=PaymentUtil.gen_serial_number(TimeUtil.format_time(time_str)),
                    status=REFUND_STATUS['REQUESTED'],
                    payment=payment,
                    datetime=TimeUtil.format_time(time_str)
                )
            else:
                raise Exception("refund's payment's status is invalid")

            payment.status = PAYMENT_STATUS['REFUND_REQUESTED']
            payment.save()

            return new_refund

    @classmethod
    def mock_on_refund_success(cls,  platform_account_id, refund, time_str):
        with transaction.atomic():
            platform_account = Account.objects.select_for_update().get(id=platform_account_id)
            merchant_account = Account.objects.select_for_update().get(
                id=refund.payment.merchant.account.id)

            payment = Payment.objects.select_for_update().get(
                serial_number=refund.payment.serial_number)
            refund = Refund.objects.select_for_update().get(serial_number=refund.serial_number)
            timestamp = TimeUtil.format_time(time_str)

            if refund.status != REFUND_STATUS['REQUESTED'] \
                    or payment.status != PAYMENT_STATUS['REFUND_REQUESTED']:
                raise Exception("mock_refund failed")

            refund.status = REFUND_STATUS['FINISHED']
            refund.save()

            payment.status = PAYMENT_STATUS['REFUND']
            payment.save()

            is_wechat_pay = True if payment.pay_channel == PAY_CHANNELS['WECHAT'] else False

            if not payment.coupon:  # 没有使用优惠券

                if is_wechat_pay:
                    merchant_account.balance = merchant_account.balance - payment.order_price
                else:
                    merchant_account.alipay_balance = merchant_account.alipay_balance - payment.order_price
                merchant_account.save()

                Transaction.objects.bulk_create([
                    Transaction(
                        content_object=refund,
                        transaction_type=TRANSACTION_TYPE['MERCHANT_REFUND'],
                        datetime=timestamp,
                        account=merchant_account,
                        amount=-payment.order_price,
                        balance_after_transaction=merchant_account.balance if is_wechat_pay else merchant_account.alipay_balance
                    ),
                    Transaction(
                        content_object=refund,
                        transaction_type=TRANSACTION_TYPE['PLATFORM_EARNING_MERCHANT_REFUND'],
                        datetime=timestamp,
                        account=platform_account,
                        amount=payment.order_price,
                        balance_after_transaction=(platform_account.balance if is_wechat_pay
                                                   else platform_account.alipay_balance) + payment.order_price
                    ),
                    Transaction(
                        content_object=refund,
                        transaction_type=TRANSACTION_TYPE['PLATFORM_REFUND'],
                        datetime=timestamp,
                        account=platform_account,
                        amount=-payment.order_price,
                        balance_after_transaction=platform_account.balance if is_wechat_pay else platform_account.alipay_balance
                    )
                ])
            else:  # 有使用优惠券的情况
                paid_price = payment.order_price - payment.coupon.discount
                total_share = payment.platform_share + payment.inviter_share + payment.originator_share

                if is_wechat_pay:
                    merchant_account.balance = merchant_account.balance - (paid_price - total_share)
                    platform_account.balance = platform_account.balance - total_share
                else:
                    merchant_account.alipay_balance = merchant_account.alipay_balance - (paid_price - total_share)
                    platform_account.alipay_balance = platform_account.alipay_balance - total_share

                merchant_account.save()
                platform_account.save()

                Transaction.objects.bulk_create([
                    Transaction(
                        content_object=refund,
                        transaction_type=TRANSACTION_TYPE['MERCHANT_REFUND'],
                        datetime=timestamp,
                        account=merchant_account,
                        amount=-(paid_price - total_share),
                        balance_after_transaction=merchant_account.balance if is_wechat_pay else merchant_account.alipay_balance
                    ),
                    Transaction(
                        content_object=refund,
                        transaction_type=TRANSACTION_TYPE['PLATFORM_EARNING_MERCHANT_REFUND'],
                        datetime=timestamp,
                        account=platform_account,
                        amount=(paid_price - total_share),
                        balance_after_transaction=(platform_account.balance if is_wechat_pay
                                                   else platform_account.alipay_balance) + paid_price
                    ),
                    Transaction(
                        content_object=refund,
                        transaction_type=TRANSACTION_TYPE['PLATFORM_REFUND'],
                        datetime=timestamp,
                        account=platform_account,
                        amount=-paid_price,
                        balance_after_transaction=platform_account.balance if is_wechat_pay else platform_account.alipay_balance
                    )
                ])

                # 设置原优惠券为销毁状态，并退给用户一张新的优惠券
                coupon = Coupon.objects.select_for_update().get(id=payment.coupon_id)
                coupon.status = COUPON_STATUS.DESTROYED
                coupon.save()

                new_coupon = Coupon.objects.create(
                    rule_id=coupon.rule_id,
                    client_id=coupon.client_id,
                    discount=coupon.discount,
                    min_charge=coupon.min_charge,
                    originator_merchant_id=coupon.originator_merchant_id,
                    status=COUPON_STATUS.NOT_USED,
                    obtain_datetime=coupon.obtain_datetime,
                )
                new_coupon.save()

    @classmethod
    def mock_withdraw(cls, account, amount, time_str, withdraw_type):
        with transaction.atomic():
            timestamp = TimeUtil.format_time(time_str)

            is_wechat_pay = True if withdraw_type == WITHDRAW_TYPE['WECHAT'] else False

            account = Account.objects.select_for_update().annotate(
                has_merchant=Exists(Merchant.objects.filter(
                    account=OuterRef('id')
                ))).annotate(
                has_marketer=Exists(Marketer.objects.filter(
                    account=OuterRef('id')
                ))).get(id=account.id)

            if (account.withdrawable_balance if is_wechat_pay else account.alipay_withdrawable_balance) < amount:
                raise Exception("withdrawable_balance less than amount")
            if account.has_merchant:
                transaction_type = TRANSACTION_TYPE.MERCHANT_WITHDRAW
            elif account.has_marketer:
                transaction_type = TRANSACTION_TYPE.MARKETER_WITHDRAW
            else:
                raise Exception("Don't know the owner of the account.")

            new_withdraw = Withdraw.objects.create(
                withdraw_type=WITHDRAW_TYPE['WECHAT'] if is_wechat_pay else WITHDRAW_TYPE['ALIPAY'],
                serial_number=PaymentUtil.gen_serial_number(timestamp),
                status=WITHDRAW_STATUS['PROCESSING'],
                datetime=timestamp,
                account=account,
                amount=amount
            )

            if is_wechat_pay:
                account.withdrawable_balance = account.withdrawable_balance - amount
                account.balance = account.balance - amount
            else:
                account.alipay_withdrawable_balance = account.alipay_withdrawable_balance - amount
                account.alipay_balance = account.alipay_balance - amount
            account.save()

            Transaction.objects.create(
                content_object=new_withdraw,
                transaction_type=transaction_type,
                datetime=timestamp,
                account=account,
                amount=-amount,
                balance_after_transaction=account.balance
            )
            new_withdraw.status = WITHDRAW_STATUS['FINISHED']
            new_withdraw.save()

    @classmethod
    def mock_on_payment_unfreeze(cls, platform_account_id, payment, time_str):
        with transaction.atomic():

            timestamp = TimeUtil.format_time(time_str)
            payment = Payment.objects.select_for_update().select_related('coupon').get(
                serial_number=payment.serial_number)

            if payment.status != PAYMENT_STATUS.FROZEN:
                raise Exception("payment's status is invalid")

            is_wechat_pay = True if payment.pay_channel == PAY_CHANNELS['WECHAT'] else False

            if payment.coupon:  # 使用了优惠券的情况
                merchant_account_id = payment.merchant.account_id
                originator_account_id = payment.coupon.originator_merchant.account_id
                inviter_account_id = payment.merchant.inviter.account_id

                accounts = Account.objects.select_for_update().filter(id__in=(
                    platform_account_id, merchant_account_id, originator_account_id,
                    inviter_account_id))

                accounts = {a.id: a for a in accounts}
                if len(accounts) != 4:
                    raise Account.DoesNotExist()

                platform_account = accounts[platform_account_id]
                merchant_account = accounts[merchant_account_id]
                originator_account = accounts[originator_account_id]
                inviter_account = accounts[inviter_account_id]

                if is_wechat_pay:
                    platform_account.balance -= payment.originator_share + payment.inviter_share
                    originator_account.balance += payment.originator_share
                    originator_account.withdrawable_balance += payment.originator_share
                    inviter_account.balance += payment.inviter_share
                    inviter_account.withdrawable_balance += payment.inviter_share
                    merchant_account.withdrawable_balance += payment.order_price - payment.coupon.discount - sum(
                        (payment.platform_share,
                         payment.originator_share,
                         payment.inviter_share))
                else:
                    platform_account.alipay_balance -= payment.originator_share + payment.inviter_share
                    originator_account.alipay_balance += payment.originator_share
                    originator_account.alipay_withdrawable_balance += payment.originator_share
                    inviter_account.alipay_balance += payment.inviter_share
                    inviter_account.alipay_withdrawable_balance += payment.inviter_share
                    merchant_account.alipay_withdrawable_balance += payment.order_price - payment.coupon.discount - sum(
                        (payment.platform_share,
                         payment.originator_share,
                         payment.inviter_share))

                platform_account.save()
                originator_account.save()
                inviter_account.save()
                merchant_account.save()

                payment.status = PAYMENT_STATUS.FINISHED
                payment.save()

                # 记录Transaction
                Transaction.objects.bulk_create([
                    Transaction(
                        content_object=payment,
                        transaction_type=TRANSACTION_TYPE['PLATFORM_EXPEND_MERCHANT_SHARE'],
                        datetime=timestamp,
                        account=platform_account,
                        amount=-payment.originator_share,
                        balance_after_transaction=(platform_account.balance if is_wechat_pay
                                                   else platform_account.alipay_balance) + payment.inviter_share
                    ),
                    Transaction(
                        content_object=payment,
                        transaction_type=TRANSACTION_TYPE['PLATFORM_EXPEND_MARKETER_SHARE'],
                        datetime=timestamp,
                        account=platform_account,
                        amount=-payment.inviter_share,
                        balance_after_transaction=platform_account.balance if is_wechat_pay else platform_account.alipay_balance
                    ),
                    Transaction(
                        content_object=payment,
                        transaction_type=TRANSACTION_TYPE['PLATFORM_EXPEND_PLATFORM_SHARE'],
                        datetime=timestamp,
                        account=platform_account,
                        amount=-payment.platform_share,
                        balance_after_transaction=(platform_account.balance if is_wechat_pay
                                                   else platform_account.alipay_balance) - payment.platform_share
                    ),
                    Transaction(
                        content_object=payment,
                        transaction_type=TRANSACTION_TYPE['PLATFORM_SHARE'],
                        datetime=timestamp,
                        account=platform_account,
                        amount=payment.platform_share,
                        balance_after_transaction=platform_account.balance if is_wechat_pay else platform_account.alipay_balance
                    ),
                    Transaction(
                        content_object=payment,
                        transaction_type=TRANSACTION_TYPE['MERCHANT_SHARE'],
                        datetime=timestamp,
                        account=originator_account,
                        amount=payment.originator_share,
                        balance_after_transaction=originator_account.balance if is_wechat_pay else originator_account.alipay_balance
                    ),
                    Transaction(
                        content_object=payment,
                        transaction_type=TRANSACTION_TYPE['MARKETER_SHARE'],
                        datetime=timestamp,
                        account=inviter_account,
                        amount=payment.inviter_share,
                        balance_after_transaction=inviter_account.balance if is_wechat_pay else inviter_account.alipay_balance
                    )
                ])

            else:  # 没有使用优惠券的情况
                merchant_account_id = payment.merchant.account_id
                merchant_account = Account.objects.select_for_update().get(id=merchant_account_id)

                payment.status = PAYMENT_STATUS.FINISHED
                payment.save()

                if is_wechat_pay:
                    merchant_account.withdrawable_balance += payment.order_price
                else:
                    merchant_account.alipay_withdrawable_balance += payment.order_price
                merchant_account.save()


import time

import streamlit as st
# import numpy as np
# import pandas as pd
# import plotly.express as px
# from PIL import Image

st.title('Streamlit 超入門')

st.write('プログレスバーの表示')
'Start!'

latest_iteration = st.empty()
bar = st.progress(0)

for i in range(100):
    latest_iteration.text(f'Iteration {i+1}')
    bar.progress(i + 1)
    time.sleep(0.1)

text = st.text_input('あなたの趣味を教えて下さい。')
condition = st.slider('あなたの今の調子は？', 0, 100, 50)

'あなたの趣味：', text
'コンディション：', condition

# if st.checkbox('Show Image'):
#     img = Image.open('alpha_stmoritz_train.jpg')
#     st.image(img, caption='train', use_column_width=True)

# option = st.selectbox(
#     'あなたが好きな数字を教えてください、',
#     list(range(1,11))
# )

# 'あなたの好きな数字は、', option, 'です。'

# df = pd.DataFrame(
#     np.random.rand(100, 2) / [50, 50] + [35.69, 139.70],
#     columns=['lat', 'lon']
# )

# #st.table(df.style.highlight_max(axis=0))
# st.map(df)

#df = [dict(Occupier="A", Start="2009-01-01", Finish="2009-02-28", Room="201"),
#      dict(Occupier="C", Start="2009-02-20", Finish="2009-05-30", Room="202"),
#      dict(Occupier="B", Start="2009-03-05", Finish="2009-04-15", Room="201"),
#      dict(Occupier="D", Start="2009-03-20", Finish="2009-05-30", Room="101")
#      ]
#fig = px.timeline(df, x_start="Start", x_end="Finish", y="Room", color="Room")

#st.plotly_chart(fig)


import logging

from common.registration_service import RegistrationService
from load_balancer import LoadBalancer
from proto.messages.meteo.meteo_messages_pb2 import RawMeteoData, RawPollutionData

logger = logging.getLogger(__name__)

class MeteoService:
    def __init__(self, load_balancer: LoadBalancer, registration_service: RegistrationService):
        logger.info("Initializing MeteoService")
        self._load_balancer = load_balancer
        self._registration_service = registration_service

    async def send_meteo_data(self, meteo_data: RawMeteoData):
        await self._load_balancer.send_meteo_data(meteo_data)

    async def send_pollution_data(self, pollution_data: RawPollutionData):
        await self._load_balancer.send_pollution_data(pollution_data)


from django.shortcuts import render
from django.contrib.auth.decorators import login_required
from notification.models import Notification
from twitter_user.models import Profile

@login_required
def notifications_view(request, handle):
    html = 'notification.html'
    notifcation_count = Notification.objects.filter(user=request.user).count()
    tweets = Notification.objects.filter(
        user=Profile.objects.get(handle=handle))

    for x in tweets:
        x.delete()

    return render(
        request,
        html,
        {'tweets': tweets, 'notifcation_count': notifcation_count}
    )