import time, datetime
import numpy as np
import os
import os.path as osp
import torch
import torchvision
import matplotlib.pyplot as plt
import torchvision.utils as vutils
import torch.nn.functional as F
import cv2
import glob
import random

from lib.utils.eval_utils import (
    batch_compute_similarity_transform_torch,
)
from lib.utils.geometry import batch_rodrigues

def inverse_normalize(tensor):
    mean = [0.485, 0.456, 0.406]
    std = [0.229, 0.224, 0.225]
    for t, m, s in zip(tensor, mean, std):
        t.mul_(s).add_(m)
    return tensor

def normalize(tensor):
    mean = [0.485, 0.456, 0.406]
    std = [0.229, 0.224, 0.225]
    for t, m, s in zip(tensor, mean, std):
        t.sub_(m).div_(s)
    return tensor

class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

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

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

def get_HHMMSS_from_second(seconds):
    HHMMSS = time.strftime("%H:%M:%S", time.gmtime(seconds))
    return HHMMSS

def save_checkpoint(state, output_dir, filename):
    """ save model parameter """

    save_path = "results/{}/save_pth".format(output_dir)
    if not osp.exists(save_path):
        os.makedirs(save_path)

    torch.save(state, osp.join(save_path, filename))

def save_img(img, output_dir, filename, epoch=None, img_path=None, test_dataset=None, test_idx=None, vflip=False, hflip=False):
    """ save image """
    if isinstance(img, torch.Tensor):
        img = img.cpu()

    if isinstance(img_path, int):
        img_path = str(img_path)
        root = "results/{}/save_output_train".format(output_dir)
        epoch = "epoch"+str(epoch)
        save_path = osp.join(root, epoch, img_path)

    elif epoch and img_path:
        dataset = img_path.split("/")[2]
        root = "results/{}/save_output_train/{}".format(output_dir, dataset)
        epoch = "epoch"+str(epoch)
        save_path = osp.join(root, epoch, ("_".join(img_path.split("/")[-2:])).replace(".jpg", ""))

    elif test_dataset:
        save_path = "results/{}/eval_data/{}_{}/{}".format(output_dir, test_dataset, epoch, test_idx)

    if not osp.exists(save_path):
        os.makedirs(save_path)

    # For jigsaw image with tile
    if len(img.shape) == 4: # tile, channel, width, height
        for tile_idx, _img_tile in enumerate(img):
            _filename = "_".join([filename.replace(".jpg", ""), str(tile_idx)]) + ".jpg"

            # _img = np.transpose(vutils.make_grid(_img_tile, normalize=True), (1, 2, 0))
            # _img = _img.numpy()
            _img = inverse_normalize(_img_tile).detach().numpy()
            _img = np.transpose(_img, (1, 2, 0))
            if vflip:
                _img = _img[::-1, :, :]
            if hflip:
                _img = _img[:, ::-1, :]
            plt.imsave(osp.join(save_path, _filename), _img)

    # For a image
    elif len(img.shape) == 3: # channel, width, height
        if isinstance(img, torch.Tensor):
            # if filename == "reconst_img.jpg":
            #     img = np.transpose(vutils.make_grid(img, normalize=True), (1, 2, 0))
            #     img = img.numpy()
            #     plt.imsave(osp.join(save_path, filename), img)
            if filename in ["detach.jpg", "rendering.jpg", "rendering_ren.jpg", "rendering_bg.jpg", "rendering1.jpg", "rendering_ren1.jpg",
                            "rendering2.jpg", "rendering_ren2.jpg", "rendering3.jpg", "rendering_ren3.jpg",
                            "rendering4.jpg", "rendering_ren4.jpg"]:
                img = np.transpose(img, (1, 2, 0))
                img = img.numpy()
                cv2.imwrite(osp.join(save_path, filename), 255*img[:, :, ::-1])
            else:
                img = inverse_normalize(img).detach().numpy()
                img = np.transpose(img, (1, 2, 0))
                plt.imsave(osp.join(save_path, filename), img)
        # if vflip:
        #     img = img[::-1, :, :]
        # if hflip:
        #     img = img[:, ::-1, :]

    elif len(img.shape) == 2: # width, height
        if isinstance(img, torch.Tensor):
            img = img[None, :, :]
            img = np.transpose(vutils.make_grid(img, normalize=True), (1, 2, 0))
            img = img.numpy()
            # img = inverse_normalize(img).detach().numpy()[:, :, None]
            # img = np.transpose(img, (1, 2, 0))
        if vflip:
            img = img[::-1, :, :]
        if hflip:
            img = img[:, ::-1, :]

        plt.imsave(osp.join(save_path, filename), img)

def save_all_img(img_dict, output_dir, epoch=None, img_path=None, test_dataset=None, test_idx=None, vflip=False, hflip=False):
    """
        img_dict keys: filename, value: image
    """
    for filename, img in img_dict.items():
        save_img(img, output_dir, filename, epoch, img_path, test_dataset, test_idx, vflip, hflip)

def save_mesh(verts, faces, output_dir, epoch, img_path=None, test_dataset=None, test_idx=None):
    """ save verts """
    filename = "mesh.obj"
    img = verts.cpu().numpy()
    faces = faces.cpu().numpy()

    if isinstance(img_path, int):
        img_path = str(img_path)
        root = "results/{}/save_output_train".format(output_dir)
        epoch = "epoch"+str(epoch)
        save_path = osp.join(root, epoch, img_path)

    elif test_dataset is not None and test_idx is not None:
        save_path = "results/{}/eval_data/{}_{}/{}".format(output_dir, test_dataset, epoch, test_idx)

    else:
        dataset = img_path.split("/")[2]
        root = "results/{}/save_output_train/{}".format(output_dir, dataset)
        epoch = "epoch"+str(epoch)
        save_path = osp.join(root, epoch, ("_".join(img_path.split("/")[-2:])).replace(".jpg", ""))

    if not osp.exists(save_path):
        os.makedirs(save_path)

    if len(verts.shape) == 2:
        with open(osp.join(save_path, filename), "w") as f:
            for verts_xyz in verts:
                f.write("v {} {} {}\n".format(verts_xyz[0], verts_xyz[1], verts_xyz[2]))

            for face in faces:
                f.write("f {} {} {}\n".format(face[0]+1, face[1]+1, face[2]+1))

def save_templates_info(test_templates, output_dir, filename):
    """ save templates
        save results information or train information """

    save_path = "results/{}/save_txt".format(output_dir)
    if not osp.exists(save_path):
        os.makedirs(save_path)

    _filename = osp.join(save_path, filename)

    with open(_filename, "w") as f:
        f.writelines("\n".join(test_templates))

def save_joints2d_img(gt_keypoints_2d, pred_keypoints_2d, output_dir, epoch, img_path):
    """
        save image of joints2d on 2d coordinate while training
    """
    line_list = [
        [ 0, 1 ],
        [ 1, 2 ],
        [ 3, 4 ],
        [ 4, 5 ],
        [ 6, 7 ],
        [ 7, 8 ],
        [ 8, 2 ],
        [ 8, 9 ],
        [ 9, 3 ],
        [ 2, 3 ],
        [ 8, 12],
        [ 9, 10],
        [12, 9 ],
        [10, 11],
        [12, 13],
    ]

    if isinstance(img_path, int):
        img_path = str(img_path)
        root = "results/{}/save_output_train".format(output_dir)
        epoch = "epoch"+str(epoch)
        save_path = osp.join(root, epoch, img_path)
    else:
        dataset = img_path.split("/")[2]
        root = "results/{}/save_output_train/{}".format(output_dir, dataset)
        epoch = "epoch"+str(epoch)
        save_path = osp.join(root, epoch, ("_".join(img_path.split("/")[-2:])).replace(".jpg", ""))

    if not osp.exists(save_path):
        os.makedirs(save_path)

    plt.figure()
    ax = plt.subplot()
    ax.invert_yaxis()

    gt_2ds = gt_keypoints_2d.clone()
    gt_2ds = gt_2ds.cpu()

    pred_2ds = pred_keypoints_2d.clone()
    pred_2ds = pred_2ds.cpu().detach()

    for joint_idx, (gt_2d, pred_2d) in enumerate(zip(gt_2ds, pred_2ds)):
        ax.scatter(gt_2d[0], gt_2d[1], marker='o', s=2, c="r")
        ax.text(gt_2d[0], gt_2d[1], joint_idx+1)
        ax.scatter(pred_2d[0], pred_2d[1], marker='o', s=2, c="b")
        ax.text(pred_2d[0], pred_2d[1], joint_idx+1)

    for start_point, end_point in line_list:
        start_point = start_point
        end_point = end_point
        ax.plot([gt_2ds[start_point][0], gt_2ds[end_point][0]], [gt_2ds[start_point][1], gt_2ds[end_point][1]], "r", linewidth=2)
        ax.plot([pred_2ds[start_point][0], pred_2ds[end_point][0]], [pred_2ds[start_point][1], pred_2ds[end_point][1]], "b", linewidth=2)

    plt.savefig(osp.join(save_path, "joints2d.jpg"))
    plt.close()

def save_joints3d_img(gt_keypoints_3d, pred_keypoints_3d, output_dir, epoch=None, img_path=None, test_dataset=None, test_idx=None):
    """
        save image of joints3d on 3d coordinate while training
    """
    line_list = [
        [ 0, 1 ],
        [ 1, 2 ],
        [ 3, 4 ],
        [ 4, 5 ],
        [ 6, 7 ],
        [ 7, 8 ],
        [ 8, 2 ],
        [ 8, 9 ],
        [ 9, 3 ],
        [ 2, 3 ],
        [ 8, 12],
        [ 9, 10],
        [12, 9 ],
        [10, 11],
        [12, 13],
    ]
    if isinstance(img_path, int):
        img_path = str(img_path)
        root = "results/{}/save_output_train".format(output_dir)
        epoch = "epoch"+str(epoch)
        save_path = osp.join(root, epoch, img_path)

    elif epoch and img_path:
        dataset = img_path.split("/")[2]
        root = "results/{}/save_output_train/{}".format(output_dir, dataset)
        epoch = "epoch"+str(epoch)
        save_path = osp.join(root, epoch, ("_".join(img_path.split("/")[-2:])).replace(".jpg", ""))

    elif test_dataset is not None and test_idx is not None:
        save_path = "results/{}/eval_data/{}_{}/{}".format(output_dir, test_dataset, epoch, test_idx)

    if not osp.exists(save_path):
        os.makedirs(save_path)

    plt.figure()
    ax = plt.subplot(1, 1, 1, projection='3d')
    gt_3ds = gt_keypoints_3d.clone()
    gt_3ds = gt_3ds.cpu()
    pred_3ds = pred_keypoints_3d.clone()
    pred_3ds = pred_3ds.cpu().detach()
    for joint_idx, (gt_3d, pred_3d) in enumerate(zip(gt_3ds, pred_3ds)):
        ax.scatter(gt_3d[0], gt_3d[1], gt_3d[2], marker='o', s=2, c="r")
        ax.text(gt_3d[0], gt_3d[1], gt_3d[2], joint_idx+1)
        ax.scatter(pred_3d[0], pred_3d[1], pred_3d[2], marker='o', s=2, c="b")
        ax.text(pred_3d[0], pred_3d[1], pred_3d[2], joint_idx+1)

    for start_point, end_point in line_list:
        start_point = start_point
        end_point = end_point
        ax.plot(
            [gt_3ds[start_point][0], gt_3ds[end_point][0]],
            [gt_3ds[start_point][1], gt_3ds[end_point][1]],
            [gt_3ds[start_point][2], gt_3ds[end_point][2]],
            color="r",
            linewidth=2)
        ax.plot(
            [pred_3ds[start_point][0], pred_3ds[end_point][0]],
            [pred_3ds[start_point][1], pred_3ds[end_point][1]],
            [pred_3ds[start_point][2], pred_3ds[end_point][2]],
            color="b",
            linewidth=2)

    plt.savefig(osp.join(save_path, "joints3d.jpg"))
    plt.close()

def get_acc(output, label):
    batch_size = output.shape[0]
    pred = torch.argmax(output, dim=1)
    correct = (pred==label).sum().item()
    acc = correct/batch_size * 100
    return acc

def spin2h36m_joint(spins, device):
    """
        Get h36m 14 joints from spin 49 joints
    """
    convert_matrix = torch.zeros((14, 49)).to(device)
    h36m_index_list = [25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38]
    for idx, h36m_index in enumerate(h36m_index_list):
        convert_matrix[idx, h36m_index] = 1
    h36m_joints = torch.matmul(convert_matrix, spins)
    return h36m_joints

def get_spherical_coords(X):
    # X is N x 3
    rad = np.linalg.norm(X, axis=1)
    # Inclination
    theta = np.arccos(X[:, 2] / rad)
    # Azimuth
    phi = np.arctan2(X[:, 1], X[:, 0])

    # Normalize both to be between [-1, 1]
    vv = (theta / np.pi) * 2 - 1
    uu = ((phi + np.pi) / (2*np.pi)) * 2 - 1
    # Return N x 2
    return np.stack([uu, vv],1)

def compute_uvsampler(verts, faces, tex_size=2):
    """
    For this mesh, pre-computes the UV coordinates for
    F x T x T points.
    Returns F x T x T x 2
    """
    alpha = np.arange(tex_size, dtype=np.float) / (tex_size-1)
    beta = np.arange(tex_size, dtype=np.float) / (tex_size-1)
    import itertools
    # Barycentric coordinate values
    coords = np.stack([p for p in itertools.product(*[alpha, beta])])
    vs = verts[faces]
    # Compute alpha, beta (this is the same order as NMR)
    v2 = vs[:, 2]
    v0v2 = vs[:, 0] - vs[:, 2]
    v1v2 = vs[:, 1] - vs[:, 2]
    # F x 3 x T*2
    samples = np.dstack([v0v2, v1v2]).dot(coords.T) + v2.reshape(-1, 3, 1)
    # F x T*2 x 3 points on the sphere
    samples = np.transpose(samples, (0, 2, 1))

    # Now convert these to uv.
    uv = get_spherical_coords(samples.reshape(-1, 3))
    # uv = uv.reshape(-1, len(coords), 2)

    uv = uv.reshape(-1, tex_size, tex_size, 2)
    return uv

def train_only_3task_network(
        HMR,
        context_encoder_net,
        discriminator,
        jigsaw_puzzle_net,
        rotation_net,

        loss_fn_BCE,
        loss_fn_MSE,
        loss_fn_CE,

        losses_CE,
        losses_DC,
        acces_JP,
        losses_JP,
        acces_ROT,
        losses_ROT,

        discriminator_optimizer,
        context_encoder_optimizer,
        jigsaw_puzzle_optimizer,
        rotation_optimizer,

        img,
        context_encoder_input,
        center_crop_img,
        jigsaw_input,
        jigsaw_order,
        rotation_img,
        rotation_idx,

        num_patch,
        ones,
        zeros,
        batch_size,
    ):
    ### Context Encoder ###
    # Update Discriminator

    feature_ce = HMR(context_encoder_input, return_only_features=True)
    feature_ce = feature_ce.reshape(-1, 2048, 1, 1)
    output_ce = context_encoder_net(feature_ce)
    output_ce_224 = context_encoder_input.clone()
    output_ce_224[:, :, 80:144, 80:144] = output_ce.clone()
    output_fake = discriminator(output_ce_224)
    output_real = discriminator(img)
    loss_BCE_fake = loss_fn_BCE(output_fake, zeros)
    loss_BCE_real = loss_fn_BCE(output_real, ones)
    loss_BCE = loss_BCE_fake + loss_BCE_real
    losses_DC.update(loss_BCE.item(), batch_size)
    discriminator_optimizer.zero_grad()
    loss_BCE.backward()
    discriminator_optimizer.step()

    # Update Decoder
    feature_ce = HMR(context_encoder_input, return_only_features=True)
    feature_ce = feature_ce.reshape(-1, 2048, 1, 1)
    output_ce = context_encoder_net(feature_ce)
    output_ce_224 = context_encoder_input.clone()
    output_ce_224[:, :, 80:144, 80:144] = output_ce.clone()
    output_fake = discriminator(output_ce_224)
    loss_BCE = loss_fn_BCE(output_fake, ones)
    loss_MSE = loss_fn_MSE(output_ce, center_crop_img)
    loss_ce = 0.001 * loss_BCE + 0.999 * loss_MSE
    losses_CE.update(loss_ce.item(), batch_size)
    context_encoder_optimizer.zero_grad()
    loss_ce.backward()
    context_encoder_optimizer.step()

    ### Jigsaw Puzzle ###
    # Update classifier
    _jigsaw_input = jigsaw_input.permute(1, 0, 2, 3, 4) # tile, batch, c, w, h
    feature_jp = list()
    for i in range(num_patch):
        feature_jp.append(HMR(_jigsaw_input[i], return_only_features=True))

    feature_jp.append(HMR(img, return_only_features=True))
    feature_jp = torch.cat(feature_jp, 1)
    output_jp = jigsaw_puzzle_net(feature_jp)
    acc_jp = get_acc(output_jp, jigsaw_order)
    acces_JP.update(acc_jp, batch_size)
    loss_jp = loss_fn_CE(output_jp, jigsaw_order)
    losses_JP.update(loss_jp.item(), batch_size)
    jigsaw_puzzle_optimizer.zero_grad()
    loss_jp.backward()
    jigsaw_puzzle_optimizer.step()

    ### Rotation ###
    # Update rotation net
    feature_rot = HMR(rotation_img, return_only_features=True)
    output_rot = rotation_net(feature_rot)
    acc_rot = get_acc(output_rot, rotation_idx)
    acces_ROT.update(acc_rot, batch_size)
    loss_rot = loss_fn_CE(output_rot, rotation_idx)
    losses_ROT.update(loss_rot, batch_size)
    rotation_optimizer.zero_grad()
    loss_rot.backward()
    rotation_optimizer.step()

    return output_ce_224

def train_hmr_using_3task(
        HMR,
        context_encoder_net,
        discriminator,
        jigsaw_puzzle_net,
        rotation_net,

        loss_fn_BCE,
        loss_fn_MSE,
        loss_fn_CE,

        losses_CE,
        acces_JP,
        losses_JP,
        acces_ROT,
        losses_ROT,
        losses_HMR_3task,

        img,
        context_encoder_input,
        center_crop_img,
        jigsaw_input,
        jigsaw_order,
        rotation_img,
        rotation_idx,

        num_patch,
        ones,
        zeros,
        batch_size,
        args,
    ):
    # loss for HMR - ce
    feature_ce = HMR(context_encoder_input, return_only_features=True)
    feature_ce = feature_ce.reshape(-1, 2048, 1, 1)
    output_ce = context_encoder_net(feature_ce)
    output_ce_224 = context_encoder_input.clone()
    output_ce_224[:, :, 80:144, 80:144] = output_ce.clone()
    output_fake = discriminator(output_ce_224)
    loss_BCE = loss_fn_BCE(output_fake, ones)
    loss_MSE = loss_fn_MSE(output_ce, center_crop_img)
    loss_ce = 0.001 * loss_BCE + 0.999 * loss_MSE
    losses_CE.update(loss_ce.item(), batch_size)

    # loss for HMR - jp
    _jigsaw_input = jigsaw_input.permute(1, 0, 2, 3, 4) # tile, batch, c, w, h
    feature_jp = list()
    for i in range(num_patch):
        feature_jp.append(HMR(_jigsaw_input[i], return_only_features=True))
    feature_jp.append(HMR(img, return_only_features=True))
    feature_jp = torch.cat(feature_jp, 1)
    output_jp = jigsaw_puzzle_net(feature_jp)
    acc_jp = get_acc(output_jp, jigsaw_order)
    acces_JP.update(acc_jp, batch_size)
    loss_jp = loss_fn_CE(output_jp, jigsaw_order)
    losses_JP.update(loss_jp.item(), batch_size)

    # loss for HMR - rot
    feature_rot = HMR(rotation_img, return_only_features=True)
    output_rot = rotation_net(feature_rot)
    acc_rot = get_acc(output_rot, rotation_idx)
    acces_ROT.update(acc_rot, batch_size)
    loss_rot = loss_fn_CE(output_rot, rotation_idx)
    losses_ROT.update(loss_rot, batch_size)

    loss_HMR = args.ce_weight * loss_ce + args.jp_weight * loss_jp + args.rot_weight * loss_rot
    loss_HMR = args.total_weight * loss_HMR
    losses_HMR_3task.update(loss_HMR.item(), batch_size)
    return loss_HMR, output_ce_224

def train_hmr_using_joints(
        HMR,

        loss_fn_keypoints,

        losses_HMR_joints3d,

        img,
        gt_keypoints_2d,
        gt_keypoints_3d,
        has_joints3d,

        joint_mapper_gt,
        batch_size,
        device,
        args,
    ):

    ### training HMR resnet update using joints info
    output = HMR(img)
    output = output[-1]

    ### calcuate loss of 2d joints ###
    pred_keypoints_2d = output["kp_2d"]

    conf = gt_keypoints_2d[:, :, -1].unsqueeze(-1).clone()
    joints2d_loss = (conf * loss_fn_keypoints(pred_keypoints_2d, gt_keypoints_2d[:, :, :-1])).mean()
    joints2d_loss = 5*joints2d_loss

    ### calcuate loss of 3d joints ###
    pred_keypoints_3d = output["kp_3d"][:, 25:, :]

    conf = gt_keypoints_3d[:, :, -1].unsqueeze(-1).clone()
    gt_keypoints_3d = gt_keypoints_3d[:, :, :-1].clone()
    gt_keypoints_3d = gt_keypoints_3d[has_joints3d==1]
    conf = conf[has_joints3d==1]
    pred_keypoints_3d = pred_keypoints_3d[has_joints3d==1]

    if len(gt_keypoints_3d) > 0:
        gt_pelvis = (gt_keypoints_3d[:, 2, :] + gt_keypoints_3d[:, 3, :]) / 2
        gt_keypoints_3d = gt_keypoints_3d - gt_pelvis[:, None, :]

        pred_pelvis = (pred_keypoints_3d[:, 2, :] + pred_keypoints_3d[:, 3, :]) / 2
        pred_keypoints_3d = pred_keypoints_3d - pred_pelvis[:, None, :]
        joints3d_loss = (conf*loss_fn_keypoints(pred_keypoints_3d, gt_keypoints_3d)).mean()

        pred_j3ds = pred_keypoints_3d[:, joint_mapper_gt, :].clone().detach()
        target_j3ds = gt_keypoints_3d[:, joint_mapper_gt, :].clone().detach()

        errors = torch.sqrt(((pred_j3ds - target_j3ds) ** 2).sum(dim=-1)).mean(dim=-1).cpu().numpy()
        S1_hat = batch_compute_similarity_transform_torch(pred_j3ds, target_j3ds)
        errors_pa = torch.sqrt(((S1_hat - target_j3ds) ** 2).sum(dim=-1)).mean(dim=-1).cpu().numpy()
        m2mm = 1000

        mpjpe = np.mean(errors) * m2mm
        pa_mpjpe = np.mean(errors_pa) * m2mm
        num_data = len(gt_keypoints_3d)
    else:
        joints3d_loss = torch.FloatTensor(1).fill_(0.).to(device).mean()
        mpjpe = np.array(0)
        pa_mpjpe = np.array(0)
        num_data = len(gt_keypoints_3d)

    total_loss = joints2d_loss

    joints3d_loss = 5*joints3d_loss
    total_loss += joints3d_loss
    losses_HMR_joints3d.update(joints3d_loss.item(), batch_size)

    total_loss *= 60
    return total_loss, mpjpe, pa_mpjpe, num_data

def train_texture_net(
        HMR,
        texture_net,
        img_renderer,

        loss_fn_MSE,
        loss_fn_mask,

        losses_texture_ori_img,
        losses_seg,
        losses_texture_total,

        texture_net_optimizer,

        img,
        black_img,

        batch_size,
        args,
        gt_mask,
        has_mask,
        train_first_stage
    ):
    output = HMR(img)
    output = output[-1]
    vertices = output['verts']
    cam = output['theta'][:, :3]

    textures = texture_net(img)
    textures = textures.expand(-1, -1, 2, 2, 2, -1)

    mask = gt_mask.clone().detach()
    mask_est, rendering = img_renderer(vertices, cam, textures)

    valid_mask = mask > 0
    valid_mask = valid_mask[:, None, :, :].type(torch.int)
    detach_images = img * valid_mask + black_img * (1-valid_mask)

    for i in range(batch_size):
        detach_images[i] = inverse_normalize(detach_images[i])

    #======================================================================================#
    # loss_texture => MSE loss (texture images with bg, original img)
    # loss_seg => MSE loss (segmentation images, target seg)
    # loss_texture_BCE => BCE loss (texture images with bg, Real(1))
    # loss_texture_total => SUM(loss_texture, loss_seg, loss_texture_BCE)
    #======================================================================================#
    loss_texture_ori_img = loss_fn_MSE(detach_images[has_mask==1], rendering[has_mask==1])
    losses_texture_ori_img.update(loss_texture_ori_img.item(), batch_size)
    loss_all = args.rendering_weight*loss_texture_ori_img

    if train_first_stage:
        texture_net_optimizer.zero_grad()
        loss_all.backward()
        texture_net_optimizer.step()
        return mask_est, detach_images, rendering, vertices
    else:
        _mask = mask_est[has_mask == 1]
        _gt_mask = gt_mask[has_mask == 1]

        if len(_mask) != 0 and len(_gt_mask) != 0:
            loss_seg = loss_fn_mask(_mask, _gt_mask)
            losses_seg.update(loss_seg.item(), batch_size)
            loss_all += args.seg_weight * loss_seg

        loss_all = args.texture_total_weight * loss_all
        losses_texture_total.update(loss_all.item(), batch_size)

        return loss_all, mask_est, detach_images, rendering, vertices

def train_hmr_using_adv_loss(
        HMR,
        texture_discriminator,
        texture_net,
        img_renderer,

        losses_disc_e,
        losses_disc,
        losses_disc_real,
        losses_disc_fake,

        img,

        batch_size,
    ):

    output = HMR(img)[-1]
    vertices = output['verts']
    cam = output['theta'][:, :3]

    textures = texture_net(img)
    textures = textures.expand(-1, -1, 2, 2, 2, -1)

    mask, rendering = img_renderer(vertices, cam, textures)

    rendering_img = rendering.clone()

    bg_idx = random.randint(1, 18)
    bg_list = glob.glob("/data/indoor_bg/train/LR/{}/color/*.png".format(bg_idx))

    cropped_bg_list = torch.zeros(batch_size, 3, 224, 224)
    for i in range(batch_size):
        random_idx = random.randint(0, len(bg_list)-1)
        bg_path = bg_list[random_idx]
        bg = cv2.imread(bg_path)
        bg_w, bg_h, _ = bg.shape
        h = w = 224
        rand_idx_w = int(np.random.randint(bg_w-w))
        rand_idx_h = int(np.random.randint(bg_h-h))
        cropped_bg = bg[rand_idx_w:rand_idx_w+w, rand_idx_h:rand_idx_h+h, :]/255.0
        cropped_bg = torch.from_numpy(cropped_bg).permute(2, 0, 1)
        cropped_bg_list[i] = cropped_bg

    cropped_bg_list = cropped_bg_list.to(rendering_img.device)
    valid_mask = mask > 0
    valid_mask = valid_mask[:, None, :, :].type(torch.int)

    rendering_img_input = valid_mask * rendering_img + (1-valid_mask) * cropped_bg_list
    rendering_bg = rendering_img_input.clone().detach()

    for i in range(batch_size):
        rendering_img_input[i] = normalize(rendering_img_input[i])

    e_disc_loss = batch_encoder_disc_l2_loss(texture_discriminator(rendering_img_input))
    losses_disc_e.update(e_disc_loss.item(), batch_size)

    fake_rendering_img_input = rendering_img_input.clone().detach()

    real = texture_discriminator(img)
    fake = texture_discriminator(fake_rendering_img_input)

    d_disc_real, d_disc_fake, d_disc_loss = batch_adv_disc_l2_loss(real, fake)
    losses_disc.update(d_disc_loss.item(), batch_size)
    losses_disc_real.update(d_disc_real.item(), batch_size)
    losses_disc_fake.update(d_disc_fake.item(), batch_size)
    return e_disc_loss, d_disc_loss, rendering_bg

def batch_encoder_disc_l2_loss(disc_value):
    '''
        Inputs:
            disc_value: N x 25
    '''
    k = disc_value.shape[0]
    return torch.sum((disc_value - 1.0) ** 2) * 1.0 / k

def batch_adv_disc_l2_loss(real_disc_value, fake_disc_value):
    '''
        Inputs:
            disc_value: N x 25
    '''
    ka = real_disc_value.shape[0]
    kb = fake_disc_value.shape[0]
    lb, la = torch.sum(fake_disc_value ** 2) / kb, torch.sum((real_disc_value - 1) ** 2) / ka
    return la, lb, la + lb


##https://towardsdatascience.com/develop-a-nlp-model-in-python-deploy-it-with-flask-step-by-step-744f3bdd7776
from flask import Flask, request, jsonify,render_template,redirect,flash
import pandas as pd
import matplotlib.pyplot as plt
#from flask_cors import CORS
from data_Preprocessing import DataPreprocessing
from vectorization import Embedding
from models import model
import os
from predict_model import predict
from dataVisualization import DataVisualization
from sklearn.metrics import confusion_matrix
from sklearn.metrics import classification_report
app = Flask(__name__, static_url_path='')
#app = Flask(__name__, static_url_path = "/static", static_folder = "static")

@app.route('/')
def home():
   return render_template('index.html')

@app.route("/upload", methods=['GET', 'POST'])
def upload_file():

    if request.method == 'POST':
        #print(request.files['file-7[]'])
       #import pdb;pdb.set_trace();
        f = request.files['file-7[]']
        #data_xls = pd.read_excel(f)
        resp = DataPreprocessing()
        data_df = resp.preprocessing(f)
        #print(data_df)
        ##Object for Vectorization
        target_names = ['Cancellation_Rescheduling','EDI_CustomerProgram','Escalation_Linedown',
                'Logistic_changes','MDL_Exclusion','NewPO_Forecast',
             'OrderEnquiry','Other','POChanges','RMA' ]
        class_vector = Embedding()
        X_train, X_test, Y_train, Y_test=class_vector.input_data(data_df)
        count_train,count_test = class_vector.Countvectorization(X_train, X_test)
        tfidf_train,tfidf_test = class_vector.TfIdfVectorization(X_train, X_test)
        ##Created Objects for models
        models=model()
        vis = DataVisualization()
        ##multinomialNB
        nb_pred_test,nb_pred_test_tfidf=models.multinomialNB(count_train,count_test,tfidf_train,tfidf_test,Y_train,Y_test)

        ##confusion matrix and classification report using CoutVectorization
        print("----NaiveBayes model Using Count Vectors----")
        print(classification_report(Y_test, nb_pred_test))
        nbcm1 = confusion_matrix(Y_test, nb_pred_test)
        fig = plt.figure(figsize=(8,6))
        ax = fig.add_subplot(111)
        vis.plot_confusion_matrix(nbcm1, classes=target_names)
        plt.savefig('/home/allu/Documents/TCSProjetcs/EmailClassification/static/images/NB_CountVector.png')

        ##confusion matrix and classification report using Tfidf
        print("------NaiveBayes model Using Tfidf -----")
        nbcm2 = confusion_matrix(Y_test,nb_pred_test_tfidf)
        fig = plt.figure(figsize=(8,6))
        ax = fig.add_subplot(111)
        vis.plot_confusion_matrix(nbcm2, classes=target_names)
        plt.savefig('/home/allu/Documents/TCSProjetcs/EmailClassification/static/images/NB_TfIdf.png')

        ##supportVectorMachine
        svmc_pred_test,svmc_pred_test_tfidf = models.supportVectorMachine(count_train,count_test,tfidf_train,tfidf_test,Y_train,Y_test)
        ##confusion matrix and classification report using CoutVectorization
        print("----SVM Using Count Vectors----")
        print(classification_report(Y_test, svmc_pred_test))
        svmcm1 = confusion_matrix(Y_test,svmc_pred_test)
        fig = plt.figure(figsize=(8,6))
        ax = fig.add_subplot(111)
        vis.plot_confusion_matrix(nbcm2, classes=target_names)
        plt.savefig('/home/allu/Documents/TCSProjetcs/EmailClassification/static/images/svmCount.png')

        ##confusion matrix and classification report using Tfidf
        print("--------SVM Tfidf------")
        print(classification_report(Y_test,svmc_pred_test_tfidf))
        svmcm1 = confusion_matrix(Y_test,svmc_pred_test_tfidf)
        fig = plt.figure(figsize=(8,6))
        ax = fig.add_subplot(111)
        vis.plot_confusion_matrix(nbcm2, classes=target_names)
        plt.savefig('/home/allu/Documents/TCSProjetcs/EmailClassification/static/images/svmTfidf.png')

        ##decisionTreeClassifier
        dtc_pred_test,dtc_pred_test_tfidf=models.decisionTreeClassifier(count_train,count_test,tfidf_train,tfidf_test,Y_train,Y_test)

        ##confusion matrix and classification report using CoutVectorization
        print("--------Decision CountVector------")
        print(classification_report(Y_test,dtc_pred_test))
        dtc1 = confusion_matrix(Y_test,dtc_pred_test)
        fig = plt.figure(figsize=(8,6))
        ax = fig.add_subplot(111)
        vis.plot_confusion_matrix(nbcm2, classes=target_names)
        plt.savefig('/home/allu/Documents/TCSProjetcs/EmailClassification/static/images/dtc_Count.png')

        ##confusion matrix and classification report using Tfidf
        print("--------Decision tfidf------")
        print(classification_report(Y_test,dtc_pred_test_tfidf))
        dtc2 = confusion_matrix(Y_test,dtc_pred_test_tfidf)
        fig = plt.figure(figsize=(8,6))
        ax = fig.add_subplot(111)
        vis.plot_confusion_matrix(nbcm2, classes=target_names)
        plt.savefig('/home/allu/Documents/TCSProjetcs/EmailClassification/static/images/dtc_Tfidf.png')

        ##randomClassifier
        random_pred_test,random_pred_test_tfidf=models.randomClassifier(count_train,count_test,tfidf_train,tfidf_test,Y_train,Y_test)

        ##confusion matrix and classification report using CoutVectorization
        print("--------RandomForest CountVector------")
        print(classification_report(Y_test,random_pred_test))
        randomclassifier1 = confusion_matrix(Y_test,random_pred_test)
        fig = plt.figure(figsize=(8,6))
        ax = fig.add_subplot(111)
        vis.plot_confusion_matrix(nbcm2, classes=target_names)
        plt.savefig('/home/allu/Documents/TCSProjetcs/EmailClassification/static/images/RF_Count.png')

        ##confusion matrix and classification report using Tfidf
        print("--------RandomForest tfidf------")
        print(classification_report(Y_test,dtc_pred_test_tfidf))
        randomclassifier2 = confusion_matrix(Y_test,random_pred_test_tfidf)
        fig = plt.figure(figsize=(8,6))
        ax = fig.add_subplot(111)
        vis.plot_confusion_matrix(nbcm2, classes=target_names)
        plt.savefig('/home/allu/Documents/TCSProjetcs/EmailClassification/static/images/RF_Tfidf.png')

        ##LogisticRegression
        logeg_test,logreg_tfidf_test= models.LogisticRegression(count_train,count_test,tfidf_train,tfidf_test,Y_train,Y_test)
        ##confusion matrix and classification report using CoutVectorization
        print("--------LogisticRegression CountVector------")
        print(classification_report(Y_test,logeg_test))
        randomclassifier1 = confusion_matrix(Y_test,logeg_test)
        fig = plt.figure(figsize=(8,6))
        ax = fig.add_subplot(111)
        vis.plot_confusion_matrix(nbcm2, classes=target_names)
        plt.savefig('/home/allu/Documents/TCSProjetcs/EmailClassification/static/images/logreg_Count.png')

        ##confusion matrix and classification report using Tfidf
        print("--------LogisticRegression tfidf------")
        print(classification_report(Y_test,logreg_tfidf_test))
        randomclassifier2 = confusion_matrix(Y_test,logreg_tfidf_test)
        fig = plt.figure(figsize=(8,6))
        ax = fig.add_subplot(111)
        vis.plot_confusion_matrix(nbcm2, classes=target_names)
        plt.savefig('/home/allu/Documents/TCSProjetcs/EmailClassification/static/images/logreg_Tfidf.png')

        ##XGBootClassification
        xgb_pred_test,xgb_pred_test_tfidf=models.XGBootClassification(count_train,count_test,tfidf_train,tfidf_test,Y_train,Y_test)

        ##confusion matrix and classification report using CoutVectorization
        print("-------- XGBootClassification CountVector------")
        print(classification_report(Y_test,xgb_pred_test))
        randomclassifier1 = confusion_matrix(Y_test,xgb_pred_test)
        fig = plt.figure(figsize=(8,6))
        ax = fig.add_subplot(111)
        vis.plot_confusion_matrix(nbcm2, classes=target_names)
        plt.savefig('/home/allu/Documents/TCSProjetcs/EmailClassification/static/images/xgb_Count.png')

        ##confusion matrix and classification report using Tfidf
        print("--------XGBootClassification  tfidf------")
        print(classification_report(Y_test,xgb_pred_test_tfidf))
        randomclassifier2 = confusion_matrix(Y_test,xgb_pred_test_tfidf)
        fig = plt.figure(figsize=(8,6))
        ax = fig.add_subplot(111)
        vis.plot_confusion_matrix(nbcm2, classes=target_names)
        plt.savefig('/home/allu/Documents/TCSProjetcs/EmailClassification/static/images/xgb_Tfidf.png')

        ##KNNCLassification
        modelknn_test, modelknn_tfidf_test = models.KNNCLassification(count_train,count_test,tfidf_train,tfidf_test,Y_train,Y_test)
        ##confusion matrix and classification report using CoutVectorization
        print("-------- KNN Classification CountVector------")
        print(classification_report(Y_test,modelknn_test))
        randomclassifier1 = confusion_matrix(Y_test,modelknn_test)
        fig = plt.figure(figsize=(8,6))
        ax = fig.add_subplot(111)
        vis.plot_confusion_matrix(nbcm2, classes=target_names)
        plt.savefig('/home/allu/Documents/TCSProjetcs/EmailClassification/static/images/knn_Count.png')

        ##confusion matrix and classification report using Tfidf
        print("--------KNN Classification  tfidf------")
        print(classification_report(Y_test,modelknn_tfidf_test))
        randomclassifier2 = confusion_matrix(Y_test,modelknn_tfidf_test)
        fig = plt.figure(figsize=(8,6))
        ax = fig.add_subplot(111)
        vis.plot_confusion_matrix(nbcm2, classes=target_names)
        plt.savefig('/home/allu/Documents/TCSProjetcs/EmailClassification/static/images/knn_Tfidf.png')

        return render_template('home.html')
        #return 'File Uploaded successfully'
        #print(data_xls)
        #return data_xls.to_html()
    return render_template('file.html')
        #return "File uploaded successfully"

@app.route("/predict", methods=['GET', 'POST'])
def predictor():
    p = predict()
    if request.method == 'POST':
        message = request.form['mail']
        data = [message]
        result = p.prediction(data)
        #result = str(result)
        #print(result)
        #print(type(result))
        return render_template('sample.html',  tables=[result.to_html(classes='data')], titles=result.columns.values)
        #return result
    return render_template('predict.html')

@app.route("/evalute")
def evalute():
   return render_template('dash.html')

@app.route("/export", methods=['GET'])
def export_records():
    return

if __name__ == "__main__":
    app.run()


from my_radial_grid import *

beta = 0.4
N = 150*(2 + 1)
ae_grid = MyAERadialGridDescriptor(beta/N, 1.0/N, N)

# Grid points in this case is
#                     a g
#            r(g) = -------,  g = 0, 1, ..., N - 1
#                   1 - b g

print("Grid parameters:")
print("a = ", ae_grid.a)
print("b = ", ae_grid.b)

# Grid points and derivative
NradialPoints = len(ae_grid)
r_g = ae_grid.r_g
dr_g = ae_grid.dr_g
for i in range(NradialPoints):
    print("{:18.10f} {:18.10f}".format(r_g[i], dr_g[i]))

#print(ae_grid.dv_g)
#print(ae_grid.default_spline_points)


# pip install requests bs4 lxml
# pip install jieba

import requests
import bs4
import jieba
import csv

stocks = set()
def prepare_stocks():
    with open('week3/Stock.csv', encoding='utf-8') as csv_file:
        csv_reader = csv.reader(csv_file)
        stock_list = list(csv_reader)
    for stock in stock_list[2:]:
        stocks.add(stock[1])

def get_text(o):
    return '' if o is None else o.text

def read_article(url):
    html = requests.get('https://www.ptt.cc' + url)
    soup = bs4.BeautifulSoup(html.text, 'lxml')

    rows = soup.select('div.push')
    for row in rows:
        content = get_text(row.select_one('.push-content'))
        issuer = get_text(row.select_one('.push-userid'))
        created = get_text(row.select_one('.push-ipdatetime'))

        tokens = jieba.lcut(content)
        keywords = {}
        for token in tokens:
            if token in stocks:
                if token in keywords:
                    keywords[token] += 1
                else:
                    keywords[token] = 1

        print(keywords)

def read_topic():
    html = requests.get('https://www.ptt.cc/bbs/Stock/index.html')
    soup = bs4.BeautifulSoup(html.text, 'lxml')

    rows = soup.select('div.r-ent')
    for row in rows:
        anchor = row.select_one('.title a')
        if anchor is not None:
            url = anchor['href']
            title = anchor.text
            count = get_text(row.select_one('.nrec'))
            issuer = get_text(row.select_one('.author'))
            created = get_text(row.select_one('.date'))

            tokens = jieba.lcut(title)
            keywords = {}
            for token in tokens:
                if token in stocks:
                    if token in keywords:
                        keywords[token] += 1
                    else:
                        keywords[token] = 1

            print(keywords)
            read_article(url)

prepare_stocks()
read_topic()


# https://swexpertacademy.com/main/code/problem/problemDetail.do?contestProbId=AV4suNtaXFEDFAUf
"""
단순한 구현 문제
사용된 코어와 사용할 코어의 개수가 최대로 사용했던 코어의 개수보다 적으면 더이상 탐색하지 않는다.
"""
dire = [[-1, 0], [1, 0], [0, -1], [0, 1]]

def DRAW(x, y, d, graph):
    tx, ty = x + dire[d][0], y + dire[d][1]
    while 0 <= tx < len(graph) and 0 <= ty < len(graph):
        if graph[tx][ty] == 1 or visit[tx][ty] == 1:
            return [0, 0]
        tx += dire[d][0]
        ty += dire[d][1]

    result = 0
    tx, ty = x + dire[d][0], y + dire[d][1]
    while 0 <= tx < len(graph) and 0 <= ty < len(graph):
        visit[tx][ty] = 1
        result += 1
        tx += dire[d][0]
        ty += dire[d][1]

    return [1, result]

def REMOVE(x, y, d, graph):
    tx, ty = x + dire[d][0], y + dire[d][1]
    while 0 <= tx < len(graph) and 0 <= ty < len(graph):
        visit[tx][ty] = 0
        tx += dire[d][0]
        ty += dire[d][1]

def BACK(depth: int, useCore : int, value : int):
    global ans, MaxCore
    # 남은 코어를 전부 사용해도 최대로 코어를 사용했을 때보다 적으면 return
    if MaxCore > useCore + (len(location) - depth): return

    if depth == len(location):
        # 사용한 코어의 개수가 같을 때 최솟값 구하기
        if useCore == MaxCore and value > ans : return
        MaxCore = useCore
        ans = value
        return

    for i in range(4):

        us = DRAW(location[depth][0], location[depth][1], i, graph)
        useCore += us[0]
        value += us[1]
        BACK(depth + 1, useCore, value)
        if us[0] == 1: REMOVE(location[depth][0], location[depth][1], i, graph)
        useCore -= us[0]
        value -= us[1]

T = int(input())
for test_case in range(1, T + 1):
    N = int(input())
    location = []
    graph = []
    visit = [[0] * N for _ in range(N)]
    for i in range(N):
        arr = list(map(int, input().split()))
        for j in range(N):
            if 0 < i < N - 1 and 0 < j < N - 1 and arr[j]:
                location.append([i, j])
        graph.append(arr)
    ans = 0x7fffffff
    MaxCore = 0
    BACK(0, 0, 0)
    print(f"#{test_case} {ans}")

BackTracking/[SW Test 샘플문제] 프로세서 연결하기.py

[SW Test 샘플문제] 프로세서 연결하기.py


import torch
import torch.nn as nn
import torch.nn.functional as F
import copy
import numpy as np
from itertools import chain

class SurgicalFineTuningBert(nn.Module):
    def __init__(
        self,
        bert_model,
    ) -> None:
        super().__init__()
        self.get_extended_attention_mask = bert_model.get_extended_attention_mask
        # copy the model

        self.opti_embedding_block = bert_model.bert.embeddings
        self.frozen_embedding_block = copy.deepcopy(self.opti_embedding_block)
        self.opti_bert_layers = bert_model.bert.encoder.layer
        self.frozen_bert_layers = copy.deepcopy(self.opti_bert_layers)
        self.opti_bert_pooler = bert_model.bert.pooler
        self.frozen_bert_pooler = copy.deepcopy(self.opti_bert_pooler)
        self.opti_bert_classifier = bert_model.classifier
        self.frozen_bert_classifier = copy.deepcopy(self.opti_bert_classifier)

        frozen_params = chain(
            self.frozen_embedding_block.parameters(),
            self.frozen_bert_layers.parameters(),
            self.frozen_bert_pooler.parameters(),
            self.frozen_bert_classifier.parameters(),
        )

        for param in frozen_params:
            param.requires_grad = False

        self.dropout = nn.Sequential(bert_model.dropout)
        if (
            "bert-small" in bert_model.name_or_path
            or "bert-med" in bert_model.name_or_path
        ):
            self.alphas = nn.Parameter(
                torch.zeros(len(bert_model.bert.encoder.layer) + 3)
            )
        else:
            self.alphas = nn.Parameter(
                torch.zeros(len(bert_model.bert.encoder.layer) + 1)
            )

    def forward(self, x):
        input_ids, attention_mask = x["input_ids"], x["attention_mask"]
        extended_attention_mask = self.get_extended_attention_mask(
            attention_mask, input_ids.size()
        )

        alpha_embeddings, alphas_layers, alpha_pooler, alpha_classifier = (
            self.alphas[0],
            self.alphas[:-2],
            self.alphas[-2],
            self.alphas[-1],
        )

        x_opti, x_frozen = self.opti_embedding_block(
            input_ids
        ), self.frozen_embedding_block(input_ids)

        a = alpha_embeddings.sigmoid()
        x = a * self.opti_embedding_block(input_ids) + (
            1 - a
        ) * self.frozen_embedding_block(input_ids)

        for i in range(len(self.opti_bert_layers)):
            a = alphas_layers[i].sigmoid()
            if i > 0:
                x_opti, x_frozen = x, x
            x = (
                a
                * self.opti_bert_layers[i](
                    x_opti, attention_mask=extended_attention_mask
                )[0]
                + (1 - a)
                * self.frozen_bert_layers[i](
                    x_frozen, attention_mask=extended_attention_mask
                )[0]
            )

        a = alpha_pooler.sigmoid()
        x = a * self.opti_bert_pooler(x) + (1 - a) * self.frozen_bert_pooler(x)
        x = self.dropout(x)

        a = alpha_classifier.sigmoid()
        x = a * self.opti_bert_classifier(x) + (1 - a) * self.frozen_bert_classifier(x)

        return x

    def forward_alphas(self, x, alphas):
        alpha_embeddings, alphas_layers, alpha_pooler, alpha_classifier = (
            alphas[0],
            alphas[:-2],
            alphas[-2],
            alphas[-1],
        )

        input_ids, attention_mask = x["input_ids"], x["attention_mask"]
        extended_attention_mask = self.get_extended_attention_mask(
            attention_mask, input_ids.size()
        )

        a = alpha_embeddings.sigmoid()
        x = a * self.opti_embedding_block(input_ids) + (
            1 - a
        ) * self.frozen_embedding_block(input_ids)

        x_opti, x_frozen = self.opti_embedding_block(
            input_ids
        ), self.frozen_embedding_block(input_ids)

        for i in range(len(self.opti_bert_layers)):
            a = alphas_layers[i].sigmoid()
            if i > 0:
                x_opti, x_frozen = x, x
            x = (
                a
                * self.opti_bert_layers[i](
                    x_opti, attention_mask=extended_attention_mask
                )[0]
                + (1 - a)
                * self.frozen_bert_layers[i](
                    x_frozen, attention_mask=extended_attention_mask
                )[0]
            )

        a = alpha_pooler.sigmoid()
        x = a * self.opti_bert_pooler(x) + (1 - a) * self.frozen_bert_pooler(x)
        x = self.dropout(x)

        a = alpha_classifier.sigmoid()
        x = a * self.opti_bert_classifier(x) + (1 - a) * self.frozen_bert_classifier(x)

        return x

    def get_alphas(self):
        return [float(a.sigmoid()) for a in self.alphas]


import os
import sys
import pdb
import torch
import numpy as np
import pickle as pkl
from PIL import Image
from random import shuffle

from torchvision import datasets, transforms

""" Template Dataset with Labels """
class XYDataset(torch.utils.data.Dataset):
    def __init__(self, x, y, **kwargs):
        self.x, self.y = x, y

        # this was to store the inverse permutation in permuted_mnist
        # so that we could 'unscramble' samples and plot them
        for name, value in kwargs.items():
            setattr(self, name, value)

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

    def __getitem__(self, idx):
        x, y = self.x[idx], self.y[idx]

        if type(x) != torch.Tensor:
            # mini_imagenet
            # we assume it's a path --> load from file
            x = self.transform(Image.open(x).convert('RGB'))
            y = torch.Tensor(1).fill_(y).long().squeeze()
        else:
            x = x.float() / 255.
            y = y.long()

        # for some reason mnist does better \in [0,1] than [-1, 1]
        if self.source == 'mnist':
            return x, y
        else:
            return (x - 0.5) * 2, y

""" Template Dataset for Continual Learning """
class CLDataLoader(object):
    def __init__(self, datasets_per_task, args, train=True):
        bs = args.batch_size if train else 64
        self.datasets = datasets_per_task

        # concat_list = []
        # for x in self.datasets:
        #     concat_list.append(x)
        #     print('loader x', x)
        self.loaders = [
                torch.utils.data.DataLoader(x, batch_size=bs, shuffle=True, drop_last=train, num_workers=0)
                for x in self.datasets ]

        # print('len(concat_list)', len(concat_list))
        # print('len(self.loaders)', len(self.loaders))

    def __getitem__(self, idx):
        return self.loaders[idx]

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

""" Split CIFAR10 into 5 tasks {{0,1}, ... {8,9}} """
def get_split_cifar10(args):
    # assert args.n_tasks in [5, 10], 'SplitCifar only works with 5 or 10 tasks'
    assert '1.' in str(torch.__version__)[:2], 'Use Pytorch 1.x!'
    args.n_tasks   = 5
    args.n_classes = 10
    args.buffer_size = args.n_tasks * args.mem_size * 2
    args.use_conv = True
    args.n_classes_per_task = 2
    args.input_size = [3, 32, 32]
    args.input_type = 'continuous'
    # because data is between [-1,1]:
    assert args.output_loss is not 'bernouilli'
    if args.output_loss == None:
        #TODO(multinomial is broken)
        #args.output_loss = 'multinomial'
        args.output_loss = 'mse'
        print('\nsetting output loss to MSE')

    # fetch MNIST
    train = datasets.CIFAR10('Data/', train=True,  download=True)
    test  = datasets.CIFAR10('Data/', train=False, download=True)

    try:
        train_x, train_y = train.data, train.targets
        test_x, test_y = test.data, test.targets
    except:
        train_x, train_y = train.train_data, train.train_labels
        test_x,  test_y  = test.test_data,   test.test_labels

    # sort according to the label
    out_train = [
        (x,y) for (x,y) in sorted(zip(train_x, train_y), key=lambda v : v[1]) ]

    out_test = [
        (x,y) for (x,y) in sorted(zip(test_x, test_y), key=lambda v : v[1]) ]

    train_x, train_y = [
            np.stack([elem[i] for elem in out_train]) for i in [0,1] ]

    test_x,  test_y  = [
            np.stack([elem[i] for elem in out_test]) for i in [0,1] ]

    train_x = torch.Tensor(train_x).permute(0, 3, 1, 2).contiguous()
    test_x  = torch.Tensor(test_x).permute(0, 3, 1, 2).contiguous()

    train_y = torch.Tensor(train_y)
    test_y  = torch.Tensor(test_y)

    # get indices of class split
    train_idx = [((train_y + i) % 10).argmax() for i in range(10)]
    train_idx = [0] + [x + 1 for x in sorted(train_idx)]

    test_idx  = [((test_y + i) % 10).argmax() for i in range(10)]
    test_idx  = [0] + [x + 1 for x in sorted(test_idx)]

    train_ds, test_ds = [], []
    skip = 10 // 5 #args.n_tasks
    for i in range(0, 10, skip):
        tr_s, tr_e = train_idx[i], train_idx[i + skip]
        te_s, te_e = test_idx[i],  test_idx[i + skip]

        train_ds += [(train_x[tr_s:tr_e], train_y[tr_s:tr_e])]
        test_ds  += [(test_x[te_s:te_e],  test_y[te_s:te_e])]

    train_ds, val_ds = make_valid_from_train(train_ds)

    train_ds = map(lambda x : XYDataset(x[0], x[1], **{'source':'cifar10'}), train_ds)
    val_ds   = map(lambda x : XYDataset(x[0], x[1], **{'source':'cifar10'}), val_ds)
    test_ds  = map(lambda x : XYDataset(x[0], x[1], **{'source':'cifar10'}), test_ds)

    return train_ds, val_ds, test_ds

def get_split_cifar100(args):
    # assert args.n_tasks in [5, 10], 'SplitCifar only works with 5 or 10 tasks'
    assert '1.' in str(torch.__version__)[:2], 'Use Pytorch 1.x!'
    args.n_tasks   = 20
    args.n_classes = 100
    args.buffer_size = args.n_tasks * args.mem_size * 5
    args.use_conv = True
    args.n_classes_per_task = 5
    args.input_size = [3, 32, 32]
    args.input_type = 'continuous'
    # because data is between [-1,1]:
    assert args.output_loss is not 'bernouilli'
    if args.output_loss == None:
        #TODO(multinomial is broken)
        #args.output_loss = 'multinomial'
        args.output_loss = 'mse'
        print('\nsetting output loss to MSE')

    # fetch MNIST
    train = datasets.CIFAR100('Data/', train=True,  download=True)
    test  = datasets.CIFAR100('Data/', train=False, download=True)

    try:
        train_x, train_y = train.data, train.targets
        test_x, test_y = test.data, test.targets
    except:
        train_x, train_y = train.train_data, train.train_labels
        test_x,  test_y  = test.test_data,   test.test_labels

    # sort according to the label
    out_train = [
        (x,y) for (x,y) in sorted(zip(train_x, train_y), key=lambda v : v[1]) ]

    out_test = [
        (x,y) for (x,y) in sorted(zip(test_x, test_y), key=lambda v : v[1]) ]

    train_x, train_y = [
            np.stack([elem[i] for elem in out_train]) for i in [0,1] ]

    test_x,  test_y  = [
            np.stack([elem[i] for elem in out_test]) for i in [0,1] ]

    train_x = torch.Tensor(train_x).permute(0, 3, 1, 2).contiguous()
    test_x  = torch.Tensor(test_x).permute(0, 3, 1, 2).contiguous()

    train_y = torch.Tensor(train_y)
    test_y  = torch.Tensor(test_y)

    # get indices of class split
    train_idx = [((train_y + i) % 100).argmax() for i in range(100)]
    train_idx = [0] + [x + 1 for x in sorted(train_idx)]

    test_idx  = [((test_y + i) % 100).argmax() for i in range(100)]
    test_idx  = [0] + [x + 1 for x in sorted(test_idx)]

    train_ds, test_ds = [], []
    skip = 100 // 20 #args.n_tasks
    for i in range(0, 100, skip):
        tr_s, tr_e = train_idx[i], train_idx[i + skip]
        te_s, te_e = test_idx[i],  test_idx[i + skip]

        train_ds += [(train_x[tr_s:tr_e], train_y[tr_s:tr_e])]
        test_ds  += [(test_x[te_s:te_e],  test_y[te_s:te_e])]

    train_ds, val_ds = make_valid_from_train(train_ds)

    train_ds = map(lambda x : XYDataset(x[0], x[1], **{'source':'cifar10'}), train_ds)
    val_ds   = map(lambda x : XYDataset(x[0], x[1], **{'source':'cifar10'}), val_ds)
    test_ds  = map(lambda x : XYDataset(x[0], x[1], **{'source':'cifar10'}), test_ds)

    return train_ds, val_ds, test_ds

def get_miniimagenet(args):

    print('loading miniimagenet dataset')
    ROOT_PATH = '/Data/Miniimagenet/'

    args.use_conv = True
    args.n_tasks   = 20
    args.n_classes = 100
    args.n_classes_per_task = 5
    args.input_size = (3, 84, 84)
    label2id = {}

    def get_data(setname):
        ds_dir = os.path.join(ROOT_PATH, setname)
        label_dirs = os.listdir(ds_dir)
        data, labels = [], []

        for label in label_dirs:
            label_dir = os.path.join(ds_dir, label)
            for image_file in os.listdir(label_dir):
                data.append(os.path.join(label_dir, image_file))
                if label not in label2id:
                    label_id = len(label2id)
                    label2id[label] = label_id
                label_id = label2id[label]
                labels.append(label_id)
        return data, labels

    transform = transforms.Compose([
        transforms.Resize(84),
        transforms.CenterCrop(84),
        transforms.ToTensor(),
    ])

    train_data, train_label = get_data('meta_train')
    valid_data, valid_label = get_data('meta_val')
    test_data,  test_label  = get_data('meta_test')

    # total of 60k examples for training, the rest for testing
    all_data  = np.array(train_data  + valid_data  + test_data)
    all_label = np.array(train_label + valid_label + test_label)

    train_ds, test_ds = [], []
    current_train, current_test = None, None

    cat = lambda x, y: np.concatenate((x, y), axis=0)

    for i in range(args.n_classes):
        class_indices = np.argwhere(all_label == i).reshape(-1)
        class_data  = all_data[class_indices]
        class_label = all_label[class_indices]
        split = int(0.8 * class_data.shape[0])

        data_train, data_test = class_data[:split], class_data[split:]
        label_train, label_test = class_label[:split], class_label[split:]

        if current_train is None:
            current_train, current_test = (data_train, label_train), (data_test, label_test)
        else:
            current_train = cat(current_train[0], data_train), cat(current_train[1], label_train)
            current_test  = cat(current_test[0],  data_test),  cat(current_test[1],  label_test)

        if i % args.n_classes_per_task == (args.n_classes_per_task  - 1):
            train_ds += [current_train]
            test_ds  += [current_test]
            current_train, current_test = None, None

    # build masks
    masks = []
    task_ids = [None for _ in range(20)]
    for task, task_data in enumerate(train_ds):
        labels = np.unique(task_data[1]) #task_data[1].unique().long()
        assert labels.shape[0] == args.n_classes_per_task
        mask = torch.zeros(args.n_classes).cuda()
        mask[labels] = 1
        masks += [mask]
        task_ids[task] = labels

    task_ids = torch.from_numpy(np.stack(task_ids)).cuda().long()
    print('task_ids', task_ids)

    train_ds, val_ds = make_valid_from_train(train_ds)
    train_ds = map(lambda x, y : XYDataset(x[0], x[1], **{'source':'cifar100', 'mask':y, 'task_ids':task_ids, 'transform':transform}), train_ds, masks)
    val_ds = map(lambda x, y: XYDataset(x[0], x[1], **{'source': 'cifar100', 'mask': y, 'task_ids': task_ids, 'transform': transform}), val_ds, masks)
    test_ds  = map(lambda x, y : XYDataset(x[0], x[1], **{'source':'cifar100', 'mask':y, 'task_ids':task_ids, 'transform':transform}), test_ds, masks)

    return train_ds, val_ds, test_ds

def make_valid_from_train(dataset, cut=0.95):
    tr_ds, val_ds = [], []
    for task_ds in dataset:
        x_t, y_t = task_ds

        # shuffle before splitting
        perm = torch.randperm(len(x_t))
        x_t, y_t = x_t[perm], y_t[perm]

        split = int(len(x_t) * cut)
        x_tr, y_tr   = x_t[:split], y_t[:split]
        x_val, y_val = x_t[split:], y_t[split:]

        tr_ds  += [(x_tr, y_tr)]
        val_ds += [(x_val, y_val)]

    return tr_ds, val_ds

class IIDDataset(torch.utils.data.Dataset):
    def __init__(self, data_loaders, seed=0):
        self.data_loader = data_loaders
        self.idx = []
        for task_id in range(len(data_loaders)):
            for i in range(len(data_loaders[task_id].dataset)):
                self.idx.append((task_id, i))
        random.Random(seed).shuffle(self.idx)

    def __getitem__(self, idx):
        task_id, instance_id = self.idx[idx]
        return self.data_loader[task_id].dataset.__getitem__(instance_id)

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


#!/usr/bin/env python

from geometry_msgs.msg import (
    PoseStamped,
    Pose,
    Point,
    Quaternion,
)

def house_coordinates(x, y, z, width, height):

    '''
    This function takes in coordinates x, y, z to determine where
    the structure will be built from. It also takes in the width and
    height of the structure you want to build and outputs coordinates
    stored in a list for each layer within the greater list. The
    coordinates are output in the correct order, for each layer the
    first brick to be picked is the middle one for odd number of brick
    layers and the right hand side middle brick for even number of
    brick layers. The  right arm will always pick up the first brick
    of each layer.
    '''

    # maximum structure is 0.8m width and 1.04m height
    # see documentaion for baxter reachable workspace http://mfg.rethinkrobotics.com/wiki/Workspace_Guidelines#tab.3DBaxter

    if y > 0.8:
        print('You have input the maximum width!')

    if z > 1.04:
        print('You have input the maximum height!')

    t = 0.06  # thickness of brick
    w = 0.09  # width of brick
    h = 0.2  # height of brick

    # bottom of table is currently (0,0,0)

    # a list which will contain all positions in order, with lists in the list for each layer
    list_of_positions = []

    # a list for the number of bricks in each layer
    layers = [5, 4, 4, 3, 3, 2, 2, 1]

    # the structure will always begin with vertical bricks as the first layer to be built

    # height determines number of layers
    # if the height is divided by the height + thickness of the brick and the remainder< than 200 (brick height= 200mm), this means that there is not enough height to build an extra vertical brick layer so the height will be an even number of layers
    if (height*1000)%((h*1000)+(t*1000)) < 200:
        # this structure's height will be divisible by h+t
        # the number of layers is therefore equal to the integer value of (height/h+t)*2
        number_of_layers = int(height/(h+t))*2
    else:
        # this structure's height will be divisible by (h+t) + 1
        # if the remainder > 200, then the height is tall enough to fit another vertical brick layer and the number of layers is equal to the integer value of (height/h+t)*2 + 1
        # number of layers is equal to the integer value of (height/h+t)*2 + 1
        number_of_layers = int(height/(h+t))*2 + 1

    # width determines which element in layers to start from
    if  int((width*10)/(h*10)) == 4:
        #start from 1st element
        layers = layers
    elif int((width*10)/(h*10)) == 3:
        #start from 3rd element (delete first 2 elements)
        del layers[0:2]
    elif int((width*10)/(h*10)) == 2:
        #start from 5th element (delete first 4 elements)
        del layers[0:4]
    elif int((width*10)/(h*10)) == 1:
        #start from 7th element (delete first 6 elements)
        del layers[0:6]

    # number_of_layers is currently defined by the height of the structure, if the width of the structure is too small that it cannot build to that height this would be an invalid number of times to run the loop for
    # when the length of the layers list is smaller than the value of the number_of_layers (when width is smaller than height) then number_of_layers is equal to instead the length of layers
    if len(layers) < number_of_layers:
        print('The width of the structure you input is too small to build this high!')
        number_of_layers = len(layers)

    #coefficients for alternating brick picking from middle for odd and even brick layers
    #coefficients include the length of bricks and the 2cm buffer between bricks
    coefficient_odd = [0, 1.1, -1.1, 2.2, -2.2]
    coefficient_even = [-0.55, 0.55, -1.65, 1.65]

    #count layers
    count_layer = 1

    # iterate through each layer
    for i in range(number_of_layers):

        # list to store coordinates for this specific layer
        layer_list = []

        # value of current layer
        current_layer = layers[i]

        # count bricks
        count_brick = 1

        # iterate through bricks in this layer
        for j in range(current_layer):

            # if layer is ODD NUMBER OF BRICKS
            # layer will be odd if the current layer %2 = 1 and even if = 0
            if current_layer%2 == 1:

                # y is horizontal coordinate
                # indexing count_brick will assign the correct coefficient
                y = coefficient_odd[count_brick-1]*h

                # z is the vertical coordinate
                # calculate how many (h+t)s are in the structure and if an extra h needs to be added for the odd layers
                # -0.03 for the gripper position from top of brick
                z= int(((count_layer)/2))*(h+t) + h*((count_layer)%2) - 0.03

                #append coordinates into layer list
                layer_list.append((x,y,z))

                #move on to next brick
                count_brick += 1

            # if layer is EVEN NUMBER OF BRICKS
            elif current_layer%2 == 0:

                # y is horizontal coordinate
                # indexing count_brick will assign the correct coefficient
                y= coefficient_even[count_brick-1]*h

                # z is the vertical coordinate
                # calculate how many (h+t)s are in the structure and if an extra h needs to be added for the odd layers
                # -0.03 for the gripper position from top of brick
                z= int(((count_layer)/2))*(h+t) + h*((count_layer)%2) - 0.03

                #append coordinates into layer list
                layer_list.append((x,y,z))

                #move on to next brick
                count_brick = count_brick + 1

        # move on to next layer
        count_layer = count_layer + 1

        #append each layer list into greater list
        list_of_positions.append(layer_list)

    print(list_of_positions)
    return list_of_positions

def posify(Coordinates, v_orientation, h_orientation):
    '''
    This function onverts all the brick postitions to poses for baxter to read
    it takes in the two types of quaternions (vertical or horizontal) and sorts
    the poses out based on the type of later (horizontal or vertical)
    '''

    #if block_poses[i][j] = int
    for i in range(len(Coordinates)):
        for j in range(len(Coordinates[i])):
            if i % 2:
                p =  Coordinates[i][j][0]
                u =  Coordinates[i][j][1]
                v =  Coordinates[i][j][2]
                Coordinates[i][j] =(Pose(
                        position=Point(x = p, y = u, z = v),
                        orientation=h_orientation))
            else:
                p =  Coordinates[i][j][0]
                u =  Coordinates[i][j][1]
                v =  Coordinates[i][j][2]
                Coordinates[i][j] =(Pose(
                        position=Point(x = p, y = u, z = v),
                        orientation=v_orientation))

    print (Coordinates)
    return Coordinates


from utility import log, timeit, Timer
import CONSTANT
from concurrent.futures import ProcessPoolExecutor
import datetime
import numpy as np
import os
import pandas as pd
import CONSTANT
import gc
from CONSTANT import LABEL,NUMERICAL_PREFIX,LABEL_CNT_SUFFIX

@timeit
def cat_Lable_Cnt_Fun(train_data, y, test_data, config):
    timer = Timer()
    cat_feature_list = [c for c in train_data if c.startswith(CONSTANT.CATEGORY_PREFIX)]
    if len(cat_feature_list) == 0: return None

    # train_data_length = len(train_data)
    train_data[LABEL] = y

    row_sp = int(np.ceil((len(train_data) + len(test_data)) / 1000000))
    col_sp = int(np.ceil(len(cat_feature_list) / 20))
    sp= row_sp * col_sp
    print(f' **** We should split it as {sp}, {col_sp}-{row_sp} sp to process! ****')
    cols_split = np.array_split(cat_feature_list, sp)
    data_list = []
    for i, cols in enumerate(cols_split):
        if len(cols) >= 1:
            pool = ProcessPoolExecutor(4)
            result_list = pool.map(cat_Lable_Cnt_Fun_sub,
                                   [[train_data[[col, LABEL]], test_data[[col]], col, config['pos_rate'], config[CONSTANT.TRAIN_LEN_OF_TRAIN_VAL]] for col in
                                    cols])

            pool.shutdown(wait=True)

            for i_data in result_list:
                if i_data is not None:
                    data_list += i_data

            print(f'{i} split successful')

    # feature_data = pd.concat(data_list, axis=1, copy=False)
    # feature_data.columns = name_list
    # timer.check("label count map done")
    # del data_list
    # gc.collect()

    test_data.drop(cat_feature_list, axis=1, inplace=True)
    cat_feature_list+=[LABEL]
    train_data.drop(cat_feature_list, axis=1, inplace=True)
    timer.check("drop")

    return data_list

def cat_Lable_Cnt_Fun_sub(params):
    train_df, test_df, c, prior, train_len_of_trainVal = params
    df_target_list = []

    col_name = NUMERICAL_PREFIX + c + LABEL_CNT_SUFFIX

    df_target_list.append(CatLabelCntClass(prior, train_len_of_trainVal).fit_transform(train_df, test_df).rename(col_name))

    if len(df_target_list) == 0:
        return None
    return df_target_list

class CatLabelCntClass:

    def __init__(self, prior, train_len_of_trainVal):
        self.prior = prior
        self.train_len_of_trainVal = train_len_of_trainVal

    def fit_transform(self, X, test_X):
        col = X.columns[0]
        label = X.columns[1]
        # CONSTANT.ROLLING_FOLD_WINDOW

        X_label_cnt_list = []

        num_per_fold = int(np.ceil(self.train_len_of_trainVal / CONSTANT.FOLD_NUM))
        index_range = np.arange(self.train_len_of_trainVal)
        for i in range(CONSTANT.FOLD_NUM - 1):
            large_split = list(index_range[ : (i+1)*num_per_fold])
            small_split = list(index_range[(i+1)*num_per_fold : (i+2)*num_per_fold])
            if len(small_split)==0: break

            label_cnt = X.iloc[large_split].groupby(col).agg({label: ['count','sum']})
            label_cnt.columns = label_cnt.columns.get_level_values(1)
            label_cnt['label_cnt'] = (label_cnt['sum'] + self.prior)/ (label_cnt['count'] + 1)
            label_cnt_result = X.iloc[small_split].reset_index().merge(label_cnt, how='left', on=col).fillna(value={'label_cnt':self.prior}).set_index('index')
            X_label_cnt_list.append(label_cnt_result['label_cnt'])
            del label_cnt, label_cnt_result

            if i == 0:
                # 处理最开始的数据
                first_roll_data = X.iloc[large_split]
                first_roll_data['label_cnt'] = self.prior
                X_label_cnt_list.append(first_roll_data['label_cnt'])
                del first_roll_data

        # 处理验证集数据
        index_range = np.arange(len(X))
        large_split = list(index_range[: self.train_len_of_trainVal])
        small_split = list(index_range[ self.train_len_of_trainVal:])
        label_cnt = X.iloc[large_split].groupby(col).agg({label: ['count', 'sum']})
        label_cnt.columns = label_cnt.columns.get_level_values(1)
        label_cnt['label_cnt'] = (label_cnt['sum'] + self.prior) / (label_cnt['count'] + 1)
        label_cnt_result = X.iloc[small_split].reset_index().merge(label_cnt, how='left', on=col).fillna(
            value={'label_cnt': self.prior}).set_index('index')
        X_label_cnt_list.append(label_cnt_result['label_cnt'])
        del label_cnt, label_cnt_result

        # 处理测试集数据
        label_cnt = X.groupby(col).agg({label: ['count', 'sum']})
        label_cnt.columns = label_cnt.columns.get_level_values(1)
        label_cnt['label_cnt'] = (label_cnt['sum'] + self.prior) / (label_cnt['count'] + 1)
        label_cnt_result = test_X.reset_index().merge(label_cnt, how='left', on=col).fillna(value={'label_cnt':self.prior}).set_index('index')
        X_label_cnt_list.append(label_cnt_result['label_cnt'])

        result = pd.concat(X_label_cnt_list, axis = 0).astype('float32')
        # result.sort_index(inplace = True)
        del label_cnt, label_cnt_result

        return result


import os
import pydicom
import numpy as np
import dicom_numpy

from utils import hidden_errors
from tf_utils import *
from pathlib import Path

def read_dicom_folder(dicom_folder, rescale=None):
    ''' Reads all .dcm files in `dicom_folder` and merges them to one volume

    Returns:
        The volume and the affine transformation from pixel indices to xyz coordinates
    '''
    dss = [pydicom.dcmread(str(dicom_folder/dcm)) for dcm in os.listdir(dicom_folder) if dcm.endswith('.dcm')]
    vol, mat = dicom_numpy.combine_slices(dss, rescale)
    return vol, dss[0]

def get_largest_dir(dirs, minsize=100):
    ''' Returns the dir with the most files from `dirs`'''
    m = max(dirs, key=lambda d: len(os.listdir(d)) if os.path.isdir(d) else 0)
    if len(os.listdir(m)) >= minsize: return m
    else: return None

def get_volume_dirs(path):
    path = Path(path)
    return list(
        filter(lambda p: p is not None,
        map(   get_largest_dir,                             # extract subdir with most files in it (highest res volume)
        map(   lambda p: list(p.iterdir()),                 # get list of actual volume directorie
        map(   lambda p: next(p.iterdir())/'Unknown Study', # cd into subfolders CQ500-CT-XX/Unknown Study/
        filter(lambda p: p.is_dir(),                        # Get all dirs, no files
        path.iterdir())))))                                 # Iterate over path directory
)

def get_volume_gen(volume_dirs, rescale=None, tf_pts=None):
    ''' Make a generator that loads volumes from a list of volume directories, `volume_dirs`.
    Returns: (volume:np.ndarray , index_to_pos_4x4:np.ndarray) '''
    def vol_gen():
        for vol_dir in volume_dirs:
            with hidden_errors():
                try:
                    vol, dcm = read_dicom_folder(vol_dir, rescale)
                    vox_scl = np.array([dicom.PixelSpacing[0], dicom.PixelSpacing[1], dicom.SliceThickness]).astype(np.float32)
                    vox_scl /= vox_scl.min()
                    vol_name = str(vol_dir.parent.parent.parent.name)
                    if tf_pts is None:
                        peaks  = get_histogram_peaks(normalized_vol)
                        tf_pts = get_trapezoid_tf_points_from_peaks(peaks)
                except dicom_numpy.DicomImportException:
                    print(f'Could not load {vol_dir}')
                    continue
            yield vol, tf_pts, vox_scl, vol_name
    return vol_gen()

__all__ = ['read_dicom_folder', 'get_largest_dir', 'get_volume_gen', 'get_volume_dirs']


#!/usr/bin/env python3

from re import M
from sys import stderr, exit

from multilanguage import Env, Lang, TALcolors
from TALinputs import TALinput
import random
import graph_connectivity_lib as gcl
from time import monotonic

# METADATA OF THIS TAL_SERVICE:
problem="graph_connectivity"
service="eval_bot_deciding_connectivity"
args_list = [
    ("goal",str),
    ("check_also_yes_certificate",bool),
    ("check_also_no_certificate",bool),
    ("code_lang",str),
    ("lang",str)
]

ENV = Env(args_list)
TAc = TALcolors(ENV)
LANG= Lang(ENV, TAc, lambda fstring: eval(f"f'{fstring}'"))

# Random n
n = random.randint(2, 1000)
# Random m
max_m = (n * (n-1))//2
if max_m > 10000: # Setting maximum
    max_m = 10000
m = random.randint(n-1, max_m)

stderr.write(f"n: {n}, m: {m}\n")

is_connected = True
seed = gcl.gen_instance_seed(connected=is_connected)

g = gcl.generate_graph(n, m, seed , TAc=TAc, LANG=LANG)

TAc.print(LANG.render_feedback("assigned-instance",f'Instance:\n n: {n}\nm: {m}\n seed: {seed}'), "yellow")
TAc.print(f'graph:\n{g.to_str()}', "white")

# Getting answer (y or n)
start = monotonic()
user_answer = input()
end = monotonic()
time = end - start

stderr.write(f"user_answer: {user_answer}\n")
stderr.write(f"is_connected: {is_connected}\n")

# Checking input validity
if (user_answer == "Y" or user_answer == "y"):
    user_answer = "yes"
if (user_answer == "N" or user_answer == "n"):
    user_answer = "no"
if (user_answer!= "yes" and user_answer!="no"):
    TAc.print(LANG.render_feedback("not-input", 'Input not valid. You can say Y,N,yes,no'),"red")
    exit(0)

# Wrong answers
if(user_answer == "yes" and is_connected == False):
    TAc.print(LANG.render_feedback("wrong-not-connected", 'WRONG, the graph is not connected'),"red")
    exit(0)
if(user_answer == "no" and is_connected == True):
    TAc.print(LANG.render_feedback("wrong-connected", 'WRONG, the graph is connected'),"red")
    exit(0)

stderr.write(f"cert - {is_connected}, {ENV['check_also_yes_certificate']}\n")
# CERTIFICATE

# yes cert
if (is_connected and ENV["check_also_yes_certificate"]):
    TAc.print(LANG.render_feedback("waiting-sp-tree",f'#? waiting for your spanning tree as routing table.\n# Format: each line two numbers separated by space. Then follow m lines, one for each edge, each with two numbers in the interval [0,n).\n# These specify the tail node and the head node of the edge, in this order.\n# Any line beggining with the \'#\' character is ignored.\n# If you prefer, you can use the \'TA_send_txt_file.py\' util here to send us the lines of a file. Just plug in the util at the \'rtal connect\' command like you do with any other bot and let the util feed in the file for you rather than acting by copy and paste yourself.'), "yellow")

    # Asking and getting sp.tree length
    TAc.print(LANG.render_feedback("waiting-sp-tree-len",'# Tell me how many rows are in your spanning tree table'), "yellow")
    start = monotonic()

    sptree_len = TALinput(int, 1, TAc=TAc)

    stderr.write(f"sp.tree len: {sptree_len}\n")
    span = gcl.Graph(n)

    has_outer_edges = True
    not_in_graph = []

    for i in range(sptree_len[0]):
        head, tail = TALinput(int, 2, TAc=TAc)
        head, tail = int(head),int(tail)

        # Checking if the inserted nodes are in the range [0, n]
        if tail >= n or head >= n or tail < 0 or head < 0:
            stderr.write(f"{head},{tail}\n")
            TAc.print(LANG.render_feedback("n-at-least-1", f'# ERROR: both ends of an edge must be nodes of the graph, i.e. integers in the range [0,{ENV["MAXN"]}.'), "red")
            exit(0)

        # check the existence of the edges (and nodes)
        if(g.check_edge(head,tail)):
            span.add_edge(head, tail)
        else:
            has_outer_edges = False
            edge = (int(head),int(tail))
            not_in_graph.append(edge)
    end = monotonic()
    time_certificate = end - start

    stderr.write(f"span\n{span.to_str()}\n")

    # check if is connect
    is_certificate_correct, not_conn = span.is_connected(True)
    is_certificate_correct = is_certificate_correct and has_outer_edges

# no cert
if (not is_connected and ENV["check_also_no_certificate"]):
    # Bipartition?
    pass

if(ENV["goal"]=="correct"):
    if(ENV["check_also_yes_certificate"] or ENV["check_also_no_certificate"]):
        if is_certificate_correct:
            TAc.print(LANG.render_feedback("correct-certificate",'Good! Your certificate is correct'),"green")
        else:
            TAc.print(LANG.render_feedback("wrong-certificate",f'WRONG! Certificate is not correct'), "red")
    else:
        TAc.print(LANG.render_feedback("right",f'Right!'), "green")

if time > 10:
    TAc.print(LANG.render_feedback("not-efficient",'Your algorithm as a whole is not very efficient, it takes more than a second\n'),"red")
else:
    TAc.print(LANG.render_feedback("efficient",'Your algorithm overall seems to be efficient!\n'),"green")
TAc.print("#end", "white")
exit(0)

'''

## With a list of instances increasing dimensions
instances = []
if ENV['goal'] == "linear":
    increase = 1.2
else: # Quadratic
    increase = 1.7

min_n = 2
max_n = 100

for i in range(8):
    n = random.randint(min_n, max_n)
    # Random m
    max_m = (n * (n-1))//2
    if max_m > 10000: # Setting maximum
        max_m = 10000
    m = random.randint(n-1, max_m)

    stderr.write(f"n: {n}, m: {m}\n")

    min_n = max_n
    max_n = int(max_n*increase)
    stderr.write(f"min_n: {min_n}, max_n: {max_n}\n")

    # Generating graph
    seed = gcl.gen_instance_seed(connected=True)
    instances.append(gcl.generate_graph(n, m, seed , TAc=TAc, LANG=LANG))

# Test
for graph in instances:
    TAc.print(LANG.render_feedback("assigned-instance",f'Instance:\n n: {n}\nm: {m}\n seed: {seed}'), "yellow")
    TAc.print(f'graph:\n{g.to_str()}', "white")

    # Getting answer (y or n)
    start = monotonic()
    user_answer = input()
    end = monotonic()
    time = end - start

    stderr.write(f"user_answer: {user_answer}\n")

    # Come sopra
'''


from django.db.models import Q
from django.shortcuts import render
from apps.news.models import News, HeadlineNews, BottomInfo

# Views function for home page of site
def index(request):
    # news part
    latest_news = News.objects.order_by('-published_date')[:3]
    headlines = HeadlineNews.objects.filter(is_publish=True)
    headline1 = headlines.get(headline="headline1")
    headline2 = headlines.get(headline="headline2")
    headline3 = headlines.get(headline="headline3")
    bottominfos = BottomInfo.objects.filter(is_publish=True)
    info1 = bottominfos.get(bottom_info="info1")
    info2 = bottominfos.get(bottom_info="info2")
    info3 = bottominfos.get(bottom_info="info3")

    context = {'latest_news': latest_news,
               'headline1': headline1,
               'headline2': headline2,
               'headline3': headline3,
               'info1': info1,
               'info2': info2,
               'info3': info3,
               }
    return render(request, 'home/index.html', context)

# View function for admin page that nested in home page
def about_us(request):
    return render(request, 'home/about_us.html')

# View function for admin page that nested in home page
def administraion(request):
    return render(request, 'home/include_admin.html')


import pandas

'''
features of papers used including:
    1.main title
    2.abstract contents
    3.author
    4.keywords
    for both main title and abstract contents, dictionary and word counts need to be attained
expected result: a dict containing text, word count, and a dictionary
for authors, authors of each paper need to be ranked by their contribution(just weight them in order)
expected result: a list of authors
'''
DATA_PATH = ""
DATA_NAME = ""

class Author:
    author_name = None
    author_publication = list()

class Paper:
    title = None
    abstract = None
    author_list = list()
    word_dict = dict()
    keywords = list()
    vectorized_keywords = list()
    label = None
    index = None
    def define(self, label, title, abstract,keywords, author_list, index, word_dict=None):
        self.label = label
        self.title = title
        self.abstract = abstract
        self.author_list = author_list
        self.keywords = keywords
        self.index = index
        self.word_dict = word_dict

    def computed_features(self, vectorized_keywords, keywords):
        self.keywords = keywords
        self.vectorized_keywords = vectorized_keywords

class Daedalus:
    corpus_word_dictionary = dict()
    corpus_author_list = dict()
    corpus_paper_list = list()

    def query_something(self, subject, target, sign):
        author_profile = Author()
        if sign:
            if any(target):
                for word in target:
                    if word == subject:
                        return 1
                return 0
            else:
                return 0
        elif not sign:
            b = any(target.values())
            if any(target.values()):
                for element in target.values():
                    if type(element) == type(author_profile):
                        if element.author_name == subject:
                            return subject
                return 0
            else:
                return 0
        else:
            return 0

    def add_paper(self, paper):
        p = Paper()
        paper_title = paper.title.split()
        word_dict = dict()
        for word in paper_title:
            t1 = self.query_something(word, word_dict.keys(), 1)
            if t1 == 1:
                word_dict[word] = word_dict[word] + 1
            if t1 == 0:
                word_dict[word] = 1
            t2 = self.query_something(word, self.corpus_word_dictionary.keys(), 1)
            if t2 == 1:
                self.corpus_word_dictionary[word] = self.corpus_word_dictionary[word] + 1
            if t2 == 0:
                self.corpus_word_dictionary[word] = 1

        paper_abstract = paper.abstract.split()
        for word in paper_abstract:
            t1 = self.query_something(word, word_dict.keys(), 1)
            if t1 == 1:
                word_dict[word] = word_dict[word] + 1
            if t1 == 0:
                word_dict[word] = 1
            t2 = self.query_something(word, self.corpus_word_dictionary.keys(), 1)
            if t2 == 1:
                self.corpus_word_dictionary[word] = self.corpus_word_dictionary[word] + 1
            if t2 == 0:
                self.corpus_word_dictionary[word] = 1
        p.define(paper.label, paper.title, paper.abstract, paper.author_list, paper.keywords, paper.index, word_dict)
        self.corpus_paper_list.append(p)
    # fetch the authors of each paper
    # rank is defined as average author rank

    def add_author(self, paper):
        paper_authors = paper.author_list.split(',')
        person = Author()
        for author in paper_authors:
            one = self.query_something(author, self.corpus_author_list, 0)
            if one is 0:
                person.author_name = author
                person.author_publication.append(paper)
                self.corpus_author_list[author] = person
            else:
                self.corpus_author_list[one].author_publication.append(paper)

    def find_paper(self, index):
        for paper in self.corpus_paper_list:
            if paper.index == index:
                return paper
        return 0
#    def whole_paper(self, paper):