import cv2

import raytracing
import open3d
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from dataset.database import parse_database_name, get_database_split, BaseDatabase
from network.field import SDFNetwork, SingleVarianceNetwork, NeRFNetwork, AppShadingNetwork, get_intersection, \
    extract_geometry, sample_pdf, MCShadingNetwork
from utils.base_utils import color_map_forward, downsample_gaussian_blur
from utils.raw_utils import linear_to_srgb

def build_imgs_info(database:BaseDatabase, img_ids):
    images = [database.get_image(img_id) for img_id in img_ids]
    poses = [database.get_pose(img_id) for img_id in img_ids]
    Ks = [database.get_K(img_id) for img_id in img_ids]

    images = np.stack(images, 0)
    images = color_map_forward(images).astype(np.float32)
    Ks = np.stack(Ks, 0).astype(np.float32)
    poses = np.stack(poses, 0).astype(np.float32)
    return {'imgs': images, 'Ks': Ks, 'poses': poses}

def imgs_info_to_torch(imgs_info, device='cpu'):
    for k, v in imgs_info.items():
        v = torch.from_numpy(v)
        if k.startswith('imgs'): v = v.permute(0,3,1,2)
        imgs_info[k] = v.to(device)
    return imgs_info

def imgs_info_slice(imgs_info, idxs):
    new_imgs_info={}
    for k, v in imgs_info.items():
        new_imgs_info[k]=v[idxs]
    return new_imgs_info

def imgs_info_to_cuda(imgs_info):
    for k, v in imgs_info.items():
        imgs_info[k]=v.cuda()
    return imgs_info

def imgs_info_downsample(imgs_info, ratio):
    b, _, h, w = imgs_info['imgs'].shape
    dh, dw = int(ratio*h), int(ratio*w)
    imgs_info_copy = {k:v for k,v in imgs_info.items()}
    imgs_info_copy['imgs'], imgs_info_copy['Ks'] = [], []
    for bi in range(b):
        img = imgs_info['imgs'][bi].cpu().numpy().transpose([1,2,0])
        img = downsample_gaussian_blur(img, ratio)
        img = cv2.resize(img, (dw,dh), interpolation=cv2.INTER_LINEAR)
        imgs_info_copy['imgs'].append(torch.from_numpy(img).permute(2,0,1))
        K = torch.from_numpy(np.diag([dw / w, dh / h, 1]).astype(np.float32)) @ imgs_info['Ks'][bi]
        imgs_info_copy['Ks'].append(K)

    imgs_info_copy['imgs'] = torch.stack(imgs_info_copy['imgs'], 0)
    imgs_info_copy['Ks'] = torch.stack(imgs_info_copy['Ks'], 0)
    return imgs_info_copy

class NeROShapeRenderer(nn.Module):
    default_cfg = {
        # standard deviation for opacity density
        'std_net': 'default',
        'std_act': 'exp',
        'inv_s_init': 0.3,
        'freeze_inv_s_step': None,

        # geometry network
        'sdf_net': 'default',
        'sdf_activation': 'none',
        'sdf_bias': 0.5,
        'sdf_n_layers': 8,
        'sdf_freq': 6,
        'sdf_d_out': 257,
        'geometry_init': True,

        # shader network
        'shader_config': {},

        # sampling strategy
        'n_samples': 64,
        'n_bg_samples': 32,
        'inf_far': 1000.0,
        'n_importance': 64,
        'up_sample_steps': 4,  # 1 for simple coarse-to-fine sampling
        'perturb': 1.0,
        'anneal_end': 50000,
        'train_ray_num': 512,
        'test_ray_num': 1024,
        'clip_sample_variance': True,

        # dataset
        'database_name': 'nerf_synthetic/lego/black_800',

        # validation
        'test_downsample_ratio': True,
        'downsample_ratio': 0.25,
        'val_geometry': False,

        # losses
        'rgb_loss': 'charbonier',
        'apply_occ_loss': True,
        'occ_loss_step': 20000,
        'occ_loss_max_pn': 2048,
        'occ_sdf_thresh': 0.01,

        "fixed_camera": False,
    }

    def __init__(self, cfg, training=True):
        super().__init__()
        self.cfg = {**self.default_cfg, **cfg}

        self.sdf_network = SDFNetwork(d_out=self.cfg['sdf_d_out'], d_in=3, d_hidden=256,
                                      n_layers=self.cfg['sdf_n_layers'],
                                      skip_in=[self.cfg['sdf_n_layers'] // 2], multires=self.cfg['sdf_freq'],
                                      bias=self.cfg['sdf_bias'], scale=1.0,
                                      geometric_init=self.cfg['geometry_init'],
                                      weight_norm=True, sdf_activation=self.cfg['sdf_activation'])

        self.deviation_network = SingleVarianceNetwork(init_val=self.cfg['inv_s_init'], activation=self.cfg['std_act'])

        # background nerf is a nerf++ model (this is outside the unit bounding sphere, so we call it outer nerf)
        self.outer_nerf = NeRFNetwork(D=8, d_in=4, d_in_view=3, W=256, multires=10, multires_view=4, output_ch=4, skips=[4], use_viewdirs=True)
        nn.init.constant_(self.outer_nerf.rgb_linear.bias, np.log(0.5))

        self.color_network = AppShadingNetwork(self.cfg['shader_config'])
        self.sdf_inter_fun = lambda x: self.sdf_network.sdf(x)

        if training:
            self._init_dataset()

    def _init_dataset(self):
        # train/test split
        self.database = parse_database_name(self.cfg['database_name'])
        self.train_ids, self.test_ids = get_database_split(self.database)
        self.train_ids = np.asarray(self.train_ids)

        self.train_imgs_info = build_imgs_info(self.database, self.train_ids)
        self.train_imgs_info = imgs_info_to_torch(self.train_imgs_info, 'cpu')
        b, _, h, w = self.train_imgs_info['imgs'].shape
        print(f'training size {h} {w} ...')
        self.train_num = len(self.train_ids)

        self.test_imgs_info = build_imgs_info(self.database, self.test_ids)
        self.test_imgs_info = imgs_info_to_torch(self.test_imgs_info, 'cpu')
        self.test_num = len(self.test_ids)

        # clean the data if we already have
        if hasattr(self, 'train_batch'):
            del self.train_batch

        self.train_batch, self.train_poses, self.tbn, _, _ = self._construct_ray_batch(self.train_imgs_info)
        self.train_poses = self.train_poses.float().cuda()

        self._shuffle_train_batch()

    def _shuffle_train_batch(self):
        self.train_batch_i = 0
        shuffle_idxs = torch.randperm(self.tbn, device='cpu')  # shuffle
        for k, v in self.train_batch.items():
            self.train_batch[k] = v[shuffle_idxs]

    def _construct_ray_batch(self, imgs_info, device='cpu'):
        imn, _, h, w = imgs_info['imgs'].shape
        coords = torch.stack(torch.meshgrid(torch.arange(h), torch.arange(w)), -1)[:, :, (1, 0)]  # h,w,2
        coords = coords.to(device)
        coords = coords.float()[None, :, :, :].repeat(imn, 1, 1, 1)  # imn,h,w,2
        coords = coords.reshape(imn, h * w, 2)
        coords = torch.cat([coords + 0.5, torch.ones(imn, h * w, 1, dtype=torch.float32, device=device)], 2)  # imn,h*w,3

        # imn,h*w,3 @ imn,3,3 => imn,h*w,3
        dirs = coords @ torch.inverse(imgs_info['Ks']).permute(0, 2, 1)
        imgs = imgs_info['imgs'].permute(0, 2, 3, 1).reshape(imn, h * w, 3)  # imn,h*w,3
        idxs = torch.arange(imn, dtype=torch.int64, device=device)[:, None, None].repeat(1, h * w, 1)  # imn,h*w,1
        poses = imgs_info['poses']  # imn,3,4

        rn = imn * h * w
        ray_batch = {
            'dirs': dirs.float().reshape(rn, 3).to(device),
            'rgbs': imgs.float().reshape(rn, 3).to(device),
            'idxs': idxs.long().reshape(rn, 1).to(device),
        }
        return ray_batch, poses, rn, h, w

    def nvs(self, pose, K, h, w):
        device = 'cuda'
        K = torch.from_numpy(K.astype(np.float32)).unsqueeze(0).to(device)
        pose = torch.from_numpy(pose.astype(np.float32)).unsqueeze(0).to(device)

        coords = torch.stack(torch.meshgrid(torch.arange(h), torch.arange(w)), -1)[:, :, (1, 0)]  # h,w,2
        coords = coords.to(device)
        coords = coords.float()[None, :, :, :].repeat(1, 1, 1, 1)  # 1,h,w,2
        coords = coords.reshape(1, h * w, 2)
        coords = torch.cat([coords + 0.5, torch.ones(1, h * w, 1, dtype=torch.float32, device=device)], 2)  # 1,h*w,3

        # 1,h*w,3 @ imn,3,3 => 1,h*w,3
        dirs = coords @ torch.inverse(K).permute(0, 2, 1)
        idxs = torch.arange(1, dtype=torch.int64, device=device)[:, None, None].repeat(1, h * w, 1)  # 1,h*w,1

        rn = h * w
        ray_batch = {
            'dirs': dirs.float().reshape(rn, 3).to(device),
            'idxs': idxs.long().reshape(rn, 1).to(device),
        }

        trn = 1024
        output_color = []
        for ri in range(0, rn, trn):
            cur_ray_batch = {}
            for k, v in ray_batch.items(): cur_ray_batch[k] = v[ri:ri + trn]

            with torch.no_grad():
                rays_o, rays_d, near, far, human_poses = self._process_ray_batch(cur_ray_batch, pose)
                cur_outputs = self.render(rays_o, rays_d, near, far, human_poses, 0, 0, is_train=False, step=300000)
                output_color.append(cur_outputs['ray_rgb'].detach().cpu().numpy())

        output_color = np.reshape(np.concatenate(output_color, 0), [h, w, 3])
        return output_color

    def get_anneal_val(self, step):
        if self.cfg['anneal_end'] < 0:
            return 1.0
        else:
            return np.min([1.0, step / self.cfg['anneal_end']])

    @staticmethod
    def near_far_from_sphere(rays_o, rays_d):
        a = torch.sum(rays_d ** 2, dim=-1, keepdim=True)
        b = 2.0 * torch.sum(rays_o * rays_d, dim=-1, keepdim=True)
        mid = 0.5 * (-b) / a
        near = mid - 1.0
        far = mid + 1.0
        near = torch.clamp(near, min=1e-3)
        return near, far

    def get_human_coordinate_poses(self, poses):
        pn = poses.shape[0]
        cam_cen = (-poses[:, :, :3].permute(0, 2, 1) @ poses[:, :, 3:])[..., 0]  # pn,3
        if self.cfg['fixed_camera']:
            pass
        else:
            cam_cen[..., 2] = 0

        Y = torch.zeros([1, 3]).expand(pn, 3)
        Y[:, 2] = -1.0
        Z = torch.clone(poses[:, 2, :3])  # pn, 3
        Z[:, 2] = 0
        Z = F.normalize(Z, dim=-1)
        X = torch.cross(Y, Z)  # pn, 3
        R = torch.stack([X, Y, Z], 1)  # pn,3,3
        t = -R @ cam_cen[:, :, None]  # pn,3,1
        return torch.cat([R, t], -1)

    def _process_ray_batch(self, ray_batch, poses):
        rays_d = ray_batch['dirs']  # rn,3
        idxs = ray_batch['idxs'][..., 0]  # rn

        rays_o = poses[:, :, :3].permute(0, 2, 1) @ -poses[:, :, 3:]  # trn,3,1
        rays_o = rays_o[idxs, :, 0]  # rn,3
        rays_d = poses[idxs, :, :3].permute(0, 2, 1) @ rays_d.unsqueeze(-1)
        rays_d = rays_d[..., 0]  # rn,3

        rays_o = rays_o
        rays_d = F.normalize(rays_d, dim=-1)
        near, far = self.near_far_from_sphere(rays_o, rays_d)

        human_poses = self.get_human_coordinate_poses(poses)
        return rays_o, rays_d, near, far, human_poses[idxs]  # rn, 3, 4

    def test_step(self, index, step,):
        target_imgs_info, target_img_ids = self.test_imgs_info, self.test_ids
        imgs_info = imgs_info_slice(target_imgs_info, torch.from_numpy(np.asarray([index], np.int64)))
        gt_depth, gt_mask = self.database.get_depth(target_img_ids[index])  # used in evaluation
        if self.cfg['test_downsample_ratio']:
            imgs_info = imgs_info_downsample(imgs_info, self.cfg['downsample_ratio'])
            h, w = gt_depth.shape
            dh, dw = int(self.cfg['downsample_ratio'] * h), int(self.cfg['downsample_ratio'] * w)
            gt_depth, gt_mask = cv2.resize(gt_depth, (dw, dh), interpolation=cv2.INTER_NEAREST), \
                cv2.resize(gt_mask.astype(np.uint8), (dw, dh), interpolation=cv2.INTER_NEAREST)
        gt_depth, gt_mask = torch.from_numpy(gt_depth), torch.from_numpy(gt_mask.astype(np.int32))
        ray_batch, input_poses, rn, h, w = self._construct_ray_batch(imgs_info)

        input_poses = input_poses.float().cuda()
        for k, v in ray_batch.items(): ray_batch[k] = v.cuda()

        trn = self.cfg['test_ray_num']
        outputs_keys = ['ray_rgb', 'gradient_error', 'normal', 'depth']
        outputs_keys += [
            'diffuse_albedo', 'diffuse_light', 'diffuse_color',
            'specular_albedo', 'specular_light', 'specular_color', 'specular_ref',
            'metallic', 'roughness', 'occ_prob', 'indirect_light', 'occ_prob_gt',
        ]
        if self.color_network.cfg['human_light']:
            outputs_keys += ['human_light']

        outputs = {k: [] for k in outputs_keys}
        for ri in range(0, rn, trn):
            cur_ray_batch= {k:v[ri:ri+trn] for k, v in ray_batch.items()}
            rays_o, rays_d, near, far, human_poses = self._process_ray_batch(cur_ray_batch, input_poses)
            cur_outputs = self.render(rays_o, rays_d, near, far, human_poses, 0, 0, is_train=False, step=step)
            for k in outputs_keys: outputs[k].append(cur_outputs[k].detach())

        for k in outputs_keys: outputs[k] = torch.cat(outputs[k], 0)
        outputs['loss_rgb'] = self.compute_rgb_loss(outputs['ray_rgb'], ray_batch['rgbs'])
        outputs['gt_rgb'] = ray_batch['rgbs'].reshape(h, w, 3)
        outputs['ray_rgb'] = outputs['ray_rgb'].reshape(h, w, 3)

        # used in evaluation
        outputs['gt_depth'] = gt_depth.unsqueeze(-1)
        outputs['gt_mask'] = gt_mask.unsqueeze(-1)

        self.zero_grad()
        return outputs

    def train_step(self, step):
        rn = self.cfg['train_ray_num']
        # fetch to gpu
        train_ray_batch = {k: v[self.train_batch_i:self.train_batch_i + rn].cuda() for k, v in self.train_batch.items()}
        self.train_batch_i += rn
        if self.train_batch_i + rn >= self.tbn: self._shuffle_train_batch()
        train_poses = self.train_poses.cuda()
        rays_o, rays_d, near, far, human_poses = self._process_ray_batch(train_ray_batch, train_poses)

        outputs = self.render(rays_o, rays_d, near, far, human_poses, -1, self.get_anneal_val(step), is_train=True, step=step)
        outputs['loss_rgb'] = self.compute_rgb_loss(outputs['ray_rgb'], train_ray_batch['rgbs'])  # ray_loss
        return outputs

    def compute_rgb_loss(self, rgb_pr, rgb_gt):
        if self.cfg['rgb_loss'] == 'l2':
            rgb_loss = torch.sum((rgb_pr - rgb_gt) ** 2, -1)
        elif self.cfg['rgb_loss'] == 'l1':
            rgb_loss = torch.sum(F.l1_loss(rgb_pr, rgb_gt, reduction='none'), -1)
        elif self.cfg['rgb_loss'] == 'smooth_l1':
            rgb_loss = torch.sum(F.smooth_l1_loss(rgb_pr, rgb_gt, reduction='none', beta=0.25), -1)
        elif self.cfg['rgb_loss'] == 'charbonier':
            epsilon = 0.001
            rgb_loss = torch.sqrt(torch.sum((rgb_gt - rgb_pr) ** 2, dim=-1) + epsilon)
        else:
            raise NotImplementedError
        return rgb_loss

    def density_activation(self, density, dists):
        return 1.0 - torch.exp(-F.softplus(density) * dists)

    def compute_density(self, points):
        points_norm = torch.norm(points, dim=-1, keepdim=True)
        points_norm = torch.clamp(points_norm, min=1e-3)
        sigma = self.outer_nerf.density(torch.cat([points / points_norm, 1.0 / points_norm], -1))[..., 0]
        return sigma

    @staticmethod
    def upsample(rays_o, rays_d, z_vals, sdf, n_importance, inv_s):
        """
        Up sampling give a fixed inv_s
        """
        batch_size, n_samples = z_vals.shape
        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]  # n_rays, n_samples, 3
        radius = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=False)
        inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)
        sdf = sdf.reshape(batch_size, n_samples)
        prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
        prev_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]
        mid_sdf = (prev_sdf + next_sdf) * 0.5
        cos_val = (next_sdf - prev_sdf) / (next_z_vals - prev_z_vals + 1e-5)

        prev_cos_val = torch.cat([torch.zeros([batch_size, 1]), cos_val[:, :-1]], dim=-1)
        cos_val = torch.stack([prev_cos_val, cos_val], dim=-1)
        cos_val, _ = torch.min(cos_val, dim=-1, keepdim=False)
        cos_val = cos_val.clip(-1e3, 0.0) * inside_sphere

        dist = (next_z_vals - prev_z_vals)
        prev_esti_sdf = mid_sdf - cos_val * dist * 0.5
        next_esti_sdf = mid_sdf + cos_val * dist * 0.5
        prev_cdf = torch.sigmoid(prev_esti_sdf * inv_s)
        next_cdf = torch.sigmoid(next_esti_sdf * inv_s)
        alpha = (prev_cdf - next_cdf + 1e-5) / (prev_cdf + 1e-5)
        weights = alpha * torch.cumprod(
            torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]

        z_samples = sample_pdf(z_vals, weights, n_importance, det=True).detach()
        return z_samples

    def cat_z_vals(self, rays_o, rays_d, z_vals, new_z_vals, sdf, last=False):
        batch_size, n_samples = z_vals.shape
        _, n_importance = new_z_vals.shape
        pts = rays_o[:, None, :] + rays_d[:, None, :] * new_z_vals[..., :, None]
        z_vals = torch.cat([z_vals, new_z_vals], dim=-1)
        z_vals, index = torch.sort(z_vals, dim=-1)

        if not last:
            new_sdf = self.sdf_network.sdf(pts.reshape(-1, 3)).reshape(batch_size, n_importance)
            sdf = torch.cat([sdf, new_sdf], dim=-1)
            xx = torch.arange(batch_size)[:, None].expand(batch_size, n_samples + n_importance).reshape(-1)
            index = index.reshape(-1)
            sdf = sdf[(xx, index)].reshape(batch_size, n_samples + n_importance)

        return z_vals, sdf

    def sample_ray(self, rays_o, rays_d, near, far, perturb):
        n_samples = self.cfg['n_samples']
        n_bg_samples = self.cfg['n_bg_samples']
        n_importance = self.cfg['n_importance']
        up_sample_steps = self.cfg['up_sample_steps']

        # sample points
        batch_size = len(rays_o)
        z_vals = torch.linspace(0.0, 1.0, n_samples)  # sn
        z_vals = near + (far - near) * z_vals[None, :]  # rn,sn
        z_vals_outside = torch.linspace(1e-3, 1.0 - 1.0 / (n_bg_samples + 1.0), n_bg_samples)

        if perturb > 0:
            t_rand = (torch.rand([batch_size, 1]) - 0.5)
            z_vals = z_vals + t_rand * 2.0 / n_samples

            mids = .5 * (z_vals_outside[..., 1:] + z_vals_outside[..., :-1])
            upper = torch.cat([mids, z_vals_outside[..., -1:]], -1)
            lower = torch.cat([z_vals_outside[..., :1], mids], -1)
            t_rand = torch.rand([batch_size, z_vals_outside.shape[-1]])
            z_vals_outside = lower[None, :] + (upper - lower)[None, :] * t_rand

        z_vals_outside = far / torch.flip(z_vals_outside, dims=[-1]) + 1.0 / n_bg_samples

        # Up sample
        with torch.no_grad():
            pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]
            sdf = self.sdf_network.sdf(pts).reshape(batch_size, n_samples)

            for i in range(up_sample_steps):
                rn, sn = z_vals.shape
                if self.cfg['clip_sample_variance']:
                    inv_s = self.deviation_network(torch.empty([1, 3])).expand(rn, sn - 1)
                    inv_s = torch.clamp(inv_s, max=64 * 2 ** i)  # prevent too large inv_s
                else:
                    inv_s = torch.ones(rn, sn - 1) * 64 * 2 ** i
                new_z_vals = self.upsample(rays_o, rays_d, z_vals, sdf, n_importance // up_sample_steps, inv_s)
                z_vals, sdf = self.cat_z_vals(rays_o, rays_d, z_vals, new_z_vals, sdf, last=(i + 1 == up_sample_steps))

        z_vals = torch.cat([z_vals, z_vals_outside], -1)
        return z_vals

    def render(self, rays_o, rays_d, near, far, human_poses, perturb_overwrite=-1, cos_anneal_ratio=0.0, is_train=True, step=None):
        """
        :param rays_o: rn,3
        :param rays_d: rn,3
        :param near:   rn,1
        :param far:    rn,1
        :param human_poses:     rn,3,4
        :param perturb_overwrite: set 0 for inference
        :param cos_anneal_ratio:
        :param is_train:
        :param step:
        :return:
        """
        perturb = self.cfg['perturb']
        if perturb_overwrite >= 0:
            perturb = perturb_overwrite
        z_vals = self.sample_ray(rays_o, rays_d, near, far, perturb)
        ret = self.render_core(rays_o, rays_d, z_vals, human_poses, cos_anneal_ratio=cos_anneal_ratio, step=step, is_train=is_train)
        return ret

    def compute_validation_info(self, z_vals, rays_o, rays_d, weights, human_poses, step):
        depth = torch.sum(weights * z_vals, -1, keepdim=True)  # rn,
        points = depth * rays_d + rays_o  # rn,3
        gradients = self.sdf_network.gradient(points)  # rn,3
        inner_mask = torch.norm(points, dim=-1, keepdim=True) <= 1.0
        outputs = {
            'depth': depth,  # rn,1
            'normal': ((F.normalize(gradients, dim=-1) + 1.0) * 0.5) * inner_mask,  # rn,3
        }

        feature_vector = self.sdf_network(points)[..., 1:]  # rn,f
        _, occ_info, inter_results = self.color_network(points, gradients, -F.normalize(rays_d, dim=-1), feature_vector, human_poses, inter_results=True, step=step)
        _, occ_prob, _ = get_intersection(self.sdf_inter_fun, self.deviation_network, points, occ_info['reflective'], sn0=128, sn1=9)  # pn,sn-1
        occ_prob_gt = torch.sum(occ_prob, dim=-1, keepdim=True)
        outputs['occ_prob_gt'] = occ_prob_gt
        for k, v in inter_results.items(): inter_results[k] = v * inner_mask
        outputs.update(inter_results)
        return outputs

    def compute_sdf_alpha(self, points, dists, dirs, cos_anneal_ratio, step):
        # points [...,3] dists [...] dirs[...,3]
        sdf_nn_output = self.sdf_network(points)
        sdf = sdf_nn_output[..., 0]
        feature_vector = sdf_nn_output[..., 1:]

        gradients = self.sdf_network.gradient(points)  # ...,3
        inv_s = self.deviation_network(points).clip(1e-6, 1e6)  # ...,1
        inv_s = inv_s[..., 0]

        if self.cfg['freeze_inv_s_step'] is not None and step < self.cfg['freeze_inv_s_step']:
            inv_s = inv_s.detach()

        true_cos = (dirs * gradients).sum(-1)  # [...]
        iter_cos = -(F.relu(-true_cos * 0.5 + 0.5) * (1.0 - cos_anneal_ratio) +
                     F.relu(-true_cos) * cos_anneal_ratio)  # always non-positive

        # Estimate signed distances at section points
        estimated_next_sdf = sdf + iter_cos * dists * 0.5
        estimated_prev_sdf = sdf - iter_cos * dists * 0.5

        prev_cdf = torch.sigmoid(estimated_prev_sdf * inv_s)
        next_cdf = torch.sigmoid(estimated_next_sdf * inv_s)

        p = prev_cdf - next_cdf
        c = prev_cdf

        alpha = ((p + 1e-5) / (c + 1e-5)).clip(0.0, 1.0)  # [...]
        return alpha, gradients, feature_vector, inv_s, sdf

    def compute_density_alpha(self, points, dists, dirs, nerf):
        norm = torch.norm(points, dim=-1, keepdim=True)
        points = torch.cat([points / norm, 1.0 / norm], -1)
        density, color = nerf(points, dirs)  # [...,1] [...,3]
        alpha = self.density_activation(density[..., 0], dists)
        color = linear_to_srgb(torch.exp(torch.clamp(color, max=5.0)))
        return alpha, color

    def compute_occ_loss(self, occ_info, points, sdf, gradients, dirs, step):
        if step < self.cfg['occ_loss_step']: return torch.zeros(1)

        occ_prob = occ_info['occ_prob']
        reflective = occ_info['reflective']

        # select a subset for occ loss
        # note we only apply occ loss on the surface
        inner_mask = torch.norm(points, dim=-1) < 0.999  # pn
        sdf_mask = torch.abs(sdf) < self.cfg['occ_sdf_thresh']
        normal_mask = torch.sum(gradients * dirs, -1) < 0  # pn
        mask = (inner_mask & normal_mask & sdf_mask)

        if torch.sum(mask) > self.cfg['occ_loss_max_pn']:
            indices = torch.nonzero(mask)[:, 0]  # npn
            idx = torch.randperm(indices.shape[0], device='cuda')  # npn
            indices = indices[idx[:self.cfg['occ_loss_max_pn']]]  # max_pn
            mask_new = torch.zeros_like(mask)
            mask_new[indices] = 1
            mask = mask_new

        if torch.sum(mask) > 0:
            inter_dist, inter_prob, inter_sdf = get_intersection(self.sdf_inter_fun, self.deviation_network, points[mask], reflective[mask], sn0=64, sn1=16)  # pn,sn-1
            occ_prob_gt = torch.sum(inter_prob, -1, keepdim=True)
            return F.l1_loss(occ_prob[mask], occ_prob_gt)
        else:
            return torch.zeros(1)

    def render_core(self, rays_o, rays_d, z_vals, human_poses, cos_anneal_ratio=0.0, step=None, is_train=True):
        batch_size, n_samples = z_vals.shape

        # section length in original space
        dists = z_vals[..., 1:] - z_vals[..., :-1]  # rn,sn-1
        dists = torch.cat([dists, dists[..., -1:]], -1)  # rn,sn
        mid_z_vals = z_vals + dists * 0.5

        points = rays_o.unsqueeze(-2) + rays_d.unsqueeze(-2) * mid_z_vals.unsqueeze(-1)
        inner_mask = torch.norm(points, dim=-1) <= 1.0
        outer_mask = ~inner_mask

        dirs = rays_d.unsqueeze(-2).expand(batch_size, n_samples, 3)
        human_poses_pt = human_poses.unsqueeze(-3).expand(batch_size, n_samples, 3, 4)
        dirs = F.normalize(dirs, dim=-1)
        alpha, sampled_color = torch.zeros(batch_size, n_samples), torch.zeros(batch_size, n_samples, 3)

        if torch.sum(outer_mask) > 0:
            alpha[outer_mask], sampled_color[outer_mask] = self.compute_density_alpha(points[outer_mask], dists[outer_mask], -dirs[outer_mask], self.outer_nerf)

        if torch.sum(inner_mask) > 0:
            alpha[inner_mask], gradients, feature_vector, inv_s, sdf = self.compute_sdf_alpha(points[inner_mask], dists[inner_mask], dirs[inner_mask], cos_anneal_ratio, step)
            sampled_color[inner_mask], occ_info = self.color_network(points[inner_mask], gradients, -dirs[inner_mask], feature_vector, human_poses_pt[inner_mask], step=step)
            # Eikonal loss
            gradient_error = (torch.linalg.norm(gradients, ord=2, dim=-1) - 1.0) ** 2
        else:
            gradient_error = torch.zeros(1)

        weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[..., :-1]  # rn,sn
        color = (sampled_color * weights[..., None]).sum(dim=1)

        outputs = {
            'ray_rgb': color,  # rn,3
            'gradient_error': gradient_error,  # rn
        }

        if torch.sum(inner_mask) > 0:
            outputs['std'] = torch.mean(1 / inv_s)
        else:
            outputs['std'] = torch.zeros(1)

        if step < 1000:
            mask = torch.norm(points, dim=-1) < 1.2
            outputs['sdf_pts'] = points[mask]
            outputs['sdf_vals'] = self.sdf_network.sdf(points[mask])[..., 0]

        if self.cfg['apply_occ_loss']:
            # occlusion loss
            if torch.sum(inner_mask) > 0:
                outputs['loss_occ'] = self.compute_occ_loss(occ_info, points[inner_mask], sdf, gradients, dirs[inner_mask], step)
            else:
                outputs['loss_occ'] = torch.zeros(1)

        if not is_train:
            outputs.update(self.compute_validation_info(z_vals, rays_o, rays_d, weights, human_poses, step))

        return outputs

    def forward(self, data):
        torch.set_default_tensor_type('torch.cuda.FloatTensor')
        is_train = 'eval' not in data
        step = data['step']

        if is_train:
            outputs = self.train_step(step)
        else:
            index = data['index']
            outputs = self.test_step(index, step=step)

            if index == 0 and self.cfg['val_geometry']:
                bbox_min = -torch.ones(3)
                bbox_max = torch.ones(3)
                vertices, triangles = extract_geometry(bbox_min, bbox_max, 128, 0, lambda x: self.sdf_network.sdf(x))
                outputs['vertices'] = vertices
                outputs['triangles'] = triangles

        torch.set_default_tensor_type('torch.FloatTensor')
        return outputs

    def predict_materials(self):
        name = self.cfg['name']
        mesh = open3d.io.read_triangle_mesh(f'data/meshes/{name}-300000.ply')
        xyz = np.asarray(mesh.vertices)
        xyz = torch.from_numpy(xyz.astype(np.float32)).cuda()
        torch.set_default_tensor_type('torch.cuda.FloatTensor')

        metallic, roughness, albedo = [], [], []
        batch_size = 8192
        for vi in range(0, xyz.shape[0], batch_size):
            feature_vectors = self.sdf_network(xyz[vi:vi + batch_size])[:, 1:]
            m, r, a = self.color_network.predict_materials(xyz[vi:vi + batch_size], feature_vectors)
            metallic.append(m.cpu().numpy())
            roughness.append(r.cpu().numpy())
            albedo.append(a.cpu().numpy())

        return {'metallic': np.concatenate(metallic, 0),
                'roughness': np.concatenate(roughness, 0),
                'albedo': np.concatenate(albedo, 0)}

class NeROMaterialRenderer(nn.Module):
    default_cfg={
        'train_ray_num': 512,
        'test_ray_num': 1024,

        'database_name': 'real/bear/raw_1024',
        'rgb_loss': 'charbonier',

        'mesh': 'data/meshes/bear_shape-300000.ply',

        'shader_cfg': {},

        'reg_mat': True,
        'reg_diffuse_light': True,
        'reg_diffuse_light_lambda': 0.1,
        'fixed_camera': False,
    }
    def __init__(self, cfg, is_train=True):
        self.cfg = {**self.default_cfg, **cfg}
        super().__init__()
        self.warned_normal = False
        self._init_geometry()
        self._init_dataset(is_train)
        self._init_shader()

    def _init_geometry(self):
        self.mesh = open3d.io.read_triangle_mesh(self.cfg['mesh'])
        self.ray_tracer = raytracing.RayTracer(np.asarray(self.mesh.vertices), np.asarray(self.mesh.triangles))

    def _init_dataset(self, is_train):
        # train/test split
        self.database = parse_database_name(self.cfg['database_name'])
        self.train_ids, self.test_ids = get_database_split(self.database, 'validation')
        self.train_ids = np.asarray(self.train_ids)

        if is_train:
            self.train_imgs_info = build_imgs_info(self.database, self.train_ids)
            self.train_imgs_info = imgs_info_to_torch(self.train_imgs_info, 'cpu')
            self.train_num = len(self.train_ids)

            self.test_imgs_info = build_imgs_info(self.database, self.test_ids)
            self.test_imgs_info = imgs_info_to_torch(self.test_imgs_info, 'cpu')
            self.test_num = len(self.test_ids)

            self.train_batch = self._construct_ray_batch(self.train_imgs_info)
            self.tbn = self.train_batch['rays_o'].shape[0]
            self._shuffle_train_batch()

    def _init_shader(self):
        self.cfg['shader_cfg']['is_real'] = self.cfg['database_name'].startswith('real')
        self.shader_network = MCShadingNetwork(self.cfg['shader_cfg'], lambda o,d: self.trace(o,d))

    def trace_in_batch(self, rays_o, rays_d, batch_size=1024**2, cpu=False):
        inters, normals, depth, hit_mask = [], [], [], []
        rn = rays_o.shape[0]
        for ri in range(0, rn, batch_size):
            inters_cur, normals_cur, depth_cur, hit_mask_cur = self.trace(rays_o[ri:ri+batch_size], rays_d[ri:ri+batch_size])
            if cpu:
                inters_cur = inters_cur.cpu()
                normals_cur = normals_cur.cpu()
                depth_cur = depth_cur.cpu()
                hit_mask_cur = hit_mask_cur.cpu()
            inters.append(inters_cur)
            normals.append(normals_cur)
            depth.append(depth_cur)
            hit_mask.append(hit_mask_cur)
        return torch.cat(inters, 0), torch.cat(normals, 0), torch.cat(depth, 0), torch.cat(hit_mask, 0)

    def trace(self, rays_o, rays_d):
        inters, normals, depth = self.ray_tracer.trace(rays_o, rays_d)
        depth = depth.reshape(*depth.shape, 1)
        normals = -normals
        normals = F.normalize(normals, dim=-1)
        if not self.warned_normal:
            print('warn!!! the normals are flipped in NeuS by default. You may flip the normal according to your mesh!')
            self.warned_normal=True
        miss_mask = depth >= 10
        hit_mask = ~miss_mask
        return inters, normals, depth, hit_mask

    def _warn_ray_tracing(self, centers):
        centers = centers.reshape([-1,3])
        distance = torch.norm(centers,dim=-1) + 1.0
        max_dist = torch.max(distance).cpu().numpy()
        if max_dist>10.0:
            print(f'warning!!! the max distance from the camera is {max_dist:.4f}, which is beyond 10.0 for the ray tracer')

    def get_human_coordinate_poses(self, poses):
        pn = poses.shape[0]
        cam_cen = (-poses[:, :, :3].permute(0,2,1) @ poses[:, :, 3:])[..., 0]  # pn,3
        if self.cfg['fixed_camera']:
            pass
        else:
            cam_cen[..., 2] = 0

        Y = torch.zeros([1, 3], device='cpu').expand(pn, 3)
        Y[:, 2] = -1.0
        Z = torch.clone(poses[:, 2, :3])  # pn, 3
        Z[:, 2] = 0
        Z = F.normalize(Z, dim=-1)
        X = torch.cross(Y, Z)  # pn, 3
        R = torch.stack([X, Y, Z], 1)  # pn,3,3
        t = -R @ cam_cen[:, :, None]  # pn,3,1
        return torch.cat([R, t], -1)

    def _construct_ray_batch(self, imgs_info, device='cpu', is_train=True):
        imn, _, h, w = imgs_info['imgs'].shape
        coords = torch.stack(torch.meshgrid(torch.arange(h),torch.arange(w)),-1)[:,:,(1,0)] # h,w,2
        coords = coords.to('cpu')
        coords = coords.float()[None,:,:,:].repeat(imn,1,1,1) # imn,h,w,2
        coords = coords.reshape(imn,h*w,2)
        coords = torch.cat([coords+0.5, torch.ones(imn, h*w, 1, dtype=torch.float32, device='cpu')], 2) # imn,h*w,3

        # imn,h*w,3 @ imn,3,3 => imn,h*w,3
        rays_d = coords @ torch.inverse(imgs_info['Ks']).permute(0, 2, 1)
        poses = imgs_info['poses']  # imn,3,4
        R, t = poses[:, :, :3], poses[:, :, 3:]
        rays_d = rays_d @ R
        rays_d = F.normalize(rays_d, dim=-1)
        rays_o = -R.permute(0,2,1) @ t # imn,3,3 @ imn,3,1
        self._warn_ray_tracing(rays_o)
        rays_o = rays_o.permute(0, 2, 1).repeat(1, h*w, 1) # imn,h*w,3
        inters, normals, depth, hit_mask = self.trace_in_batch(rays_o.reshape(-1,3), rays_d.reshape(-1,3), cpu=True) # imn
        inters, normals, depth, hit_mask = inters.reshape(imn,h*w,3), normals.reshape(imn,h*w,3), depth.reshape(imn,h*w,1), hit_mask.reshape(imn, h*w)

        human_poses = self.get_human_coordinate_poses(poses) # imn,3,4
        human_poses = human_poses.unsqueeze(1).repeat(1,h*w,1,1) # imn,h*w,3,4
        rgb = imgs_info['imgs'].reshape(imn,3,h*w).permute(0,2,1) # imn,h*w,3

        if is_train:
            ray_batch={
                'rays_o': rays_o[hit_mask].to(device),
                'rays_d': rays_d[hit_mask].to(device),
                'inters': inters[hit_mask].to(device),
                'normals': normals[hit_mask].to(device),
                'depth': depth[hit_mask].to(device),
                'human_poses': human_poses[hit_mask].to(device),
                'rgb': rgb[hit_mask].to(device),
            }
        else:
            assert imn==1
            ray_batch={
                'rays_o': rays_o[0].to(device),
                'rays_d': rays_d[0].to(device),
                'inters': inters[0].to(device),
                'normals': normals[0].to(device),
                'depth': depth[0].to(device),
                'human_poses': human_poses[0].to(device),
                'rgb': rgb[0].to(device),
                'hit_mask': hit_mask[0].to(device),
            }
        return ray_batch

    def _shuffle_train_batch(self):
        self.train_batch_i = 0
        shuffle_idxs = torch.randperm(self.tbn, device='cpu') # shuffle
        for k, v in self.train_batch.items():
            self.train_batch[k] = v[shuffle_idxs]

    def shade(self, pts, view_dirs, normals, human_poses, is_train, step=None):
        rgb_pr, outputs = self.shader_network(pts, view_dirs, normals, human_poses, step, is_train)
        outputs['rgb_pr'] = rgb_pr
        return outputs

    def compute_rgb_loss(self, rgb_pr, rgb_gt):
        if self.cfg['rgb_loss'] == 'l1':
            rgb_loss = torch.sum(F.l1_loss(rgb_pr, rgb_gt, reduction='none'), -1)
        elif self.cfg['rgb_loss'] == 'charbonier':
            epsilon = 0.001
            rgb_loss = torch.sqrt(torch.sum((rgb_gt - rgb_pr) ** 2, dim=-1) + epsilon)
        else:
            raise NotImplementedError
        return rgb_loss

    def compute_diffuse_light_regularization(self, diffuse_lights):
        diffuse_white_reg = torch.sum(torch.abs(diffuse_lights - torch.mean(diffuse_lights, dim=-1, keepdim=True)), dim=-1)
        return diffuse_white_reg * self.cfg['reg_diffuse_light_lambda']

    def train_step(self, step):
        rn = self.cfg['train_ray_num']
        pts = self.train_batch['inters'][self.train_batch_i:self.train_batch_i + rn].cuda()
        view_dirs = -self.train_batch['rays_d'][self.train_batch_i:self.train_batch_i + rn].cuda() # view directions are opposite to ray directions
        normals = self.train_batch['normals'][self.train_batch_i:self.train_batch_i + rn].cuda()
        rgb_gt = self.train_batch['rgb'][self.train_batch_i:self.train_batch_i + rn].cuda()
        human_poses = self.train_batch['human_poses'][self.train_batch_i:self.train_batch_i + rn].cuda()

        shade_outputs = self.shade(pts, view_dirs, normals, human_poses, True, step)
        shade_outputs['rgb_gt'] = rgb_gt
        shade_outputs['loss_rgb'] = self.compute_rgb_loss(shade_outputs['rgb_pr'],shade_outputs['rgb_gt'])
        if self.cfg['reg_mat']:
            shade_outputs['loss_mat_reg'] = self.shader_network.material_regularization(
                pts, normals, shade_outputs['metallic'], shade_outputs['roughness'], shade_outputs['albedo'], step)
        if self.cfg['reg_diffuse_light']:
            shade_outputs['loss_diffuse_light'] = self.compute_diffuse_light_regularization(shade_outputs['diffuse_light'])

        self.train_batch_i += rn
        if self.train_batch_i + rn >= self.tbn: self._shuffle_train_batch()
        return shade_outputs

    def test_step(self, index):
        test_imgs_info = imgs_info_slice(self.test_imgs_info, torch.from_numpy(np.asarray([index],np.int64)))
        _, _, h, w = test_imgs_info['imgs'].shape
        ray_batch = self._construct_ray_batch(test_imgs_info, 'cuda', False)
        trn = self.cfg['test_ray_num']

        output_keys = {'rgb_gt':3, 'rgb_pr':3, 'specular_light':3, 'specular_color':3, 'diffuse_light':3, 'diffuse_color':3, 'albedo':3, 'metallic':1, 'roughness':1}
        outputs = {k:[] for k in output_keys.keys()}
        rn = ray_batch['rays_o'].shape[0]
        for ri in range(0, rn, trn):
            hit_mask = ray_batch['hit_mask'][ri:ri+trn]
            outputs_cur = {k: torch.zeros(hit_mask.shape[0], d) for k, d in output_keys.items()}
            if torch.sum(hit_mask)>0:
                pts = ray_batch['inters'][ri:ri+trn][hit_mask]
                view_dirs = -ray_batch['rays_d'][ri:ri+trn][hit_mask]
                normals = ray_batch['normals'][ri:ri+trn][hit_mask]
                rgb_gt = ray_batch['rgb'][ri:ri+trn][hit_mask]
                human_poses = ray_batch['human_poses'][ri:ri+trn][hit_mask]

                shade_outputs = self.shade(pts, view_dirs, normals, human_poses, False)

                outputs_cur['rgb_pr'][hit_mask] = shade_outputs['rgb_pr']
                outputs_cur['rgb_gt'][hit_mask] = rgb_gt
                outputs_cur['specular_light'][hit_mask] = shade_outputs['specular_light']
                outputs_cur['specular_color'][hit_mask] = shade_outputs['specular_color']
                outputs_cur['diffuse_color'][hit_mask] = shade_outputs['diffuse_color']
                outputs_cur['diffuse_light'][hit_mask] = shade_outputs['diffuse_light']
                outputs_cur['albedo'][hit_mask] = shade_outputs['albedo']
                outputs_cur['metallic'][hit_mask] = shade_outputs['metallic']
                outputs_cur['roughness'][hit_mask] = torch.sqrt(shade_outputs['roughness']) # note: we assume predictions are roughness squared

            for k in output_keys.keys():
                outputs[k].append(outputs_cur[k])

        for k in output_keys.keys():
            outputs[k] = torch.cat(outputs[k], 0).reshape(h, w, -1)

        return outputs

    def forward(self, data):
        torch.set_default_tensor_type('torch.cuda.FloatTensor')
        is_train = 'eval' not in data
        step = data['step']

        if is_train:
            outputs = self.train_step(step)
        else:
            index = data['index']
            outputs = self.test_step(index)

        torch.set_default_tensor_type('torch.FloatTensor')
        return outputs

    def predict_materials(self, batch_size=8192):
        verts = torch.from_numpy(np.asarray(self.mesh.vertices,np.float32)).cuda().float()
        metallic, roughness, albedo = [], [], []
        for vi in range(0, verts.shape[0], batch_size):
            m, r, a = self.shader_network.predict_materials(verts[vi:vi+batch_size])
            r = torch.sqrt(torch.clamp(r, min=1e-7)) # note: we assume predictions are squared roughness!!!
            metallic.append(m.cpu().numpy())
            roughness.append(r.cpu().numpy())
            albedo.append(a.cpu().numpy())

        return {'metallic': np.concatenate(metallic, 0),
                'roughness': np.concatenate(roughness, 0),
                'albedo': np.concatenate(albedo, 0)}

name2renderer={
    'shape': NeROShapeRenderer,
    'material': NeROMaterialRenderer,
}


import time
import random
import os
import re
import pandas as pd
from pandas.errors import EmptyDataError, ParserError
from selenium.webdriver.firefox.options import Options
from selenium import webdriver

def rand_delay(low, high):
    time.sleep(random.uniform(low, high))

def get_file_list(_dir, suffix=None):
    files = []
    for dirpath, dirnames, filenames in os.walk(_dir):
        for filename in filenames:
            files.append(os.path.join(dirpath, filename))

    if suffix:
        files = [f for f in files if f.endswith(suffix)]
    return files

def get_latest_file(root):
    files = get_file_list(root)
    files = [f for f in files if not re.search(r'_\d+$', os.path.splitext(f)[0])]
    latest_file = max(files, key=os.path.getctime)
    return latest_file

def get_latest_session_index(root):
    dirs = os.listdir(root)
    dirs = [os.path.join(root, _dir) for _dir in dirs if re.search(r'session\d+$', _dir, re.IGNORECASE)]
    try:
        latest_dir = max(dirs, key=os.path.getctime)
        latest_index = int(latest_dir.split('\\')[-1].replace('session', ''))
        return latest_index
    except ValueError:
        return 0

def concat_data(in_files):
    dfs = []
    for file in in_files:
        try:
            try:
                df = pd.read_csv(file)
            except (ParserError, UnicodeDecodeError):
                df = pd.read_excel(file, engine='openpyxl')
            dfs.append(df)
        except EmptyDataError:
            print(f'"{file}" is empty')
    combined_data = pd.concat(dfs, join='inner', ignore_index=True)
    # combined_data.drop_duplicates(inplace=True)
    return combined_data

def setup_browser(driver_path, download_dir=None):
    profile = webdriver.FirefoxProfile()
    profile.set_preference("browser.download.folderList", 2)
    profile.set_preference("browser.download.manager.showWhenStarting", False)

    if download_dir:
        profile.set_preference("browser.download.dir", download_dir)
    profile.set_preference("browser.helperApps.neverAsk.saveToDisk", "text/csv")

    options = Options()
    options.headless = True
    browser = webdriver.Firefox(executable_path=driver_path, firefox_profile=profile, options=options)
    browser.set_window_size(1920, 1080)
    browser.set_page_load_timeout(600)
    return browser


import pandas as pd
from matplotlib import pyplot as plt
import seaborn as sns

df = pd.read_csv("train.csv") #import data from train.csv as dataframe

corr = df.Sex.str.get_dummies(sep=' ').corrwith(df.Survived/df.Survived.max()) #determines correlation between survived and sex
print(corr)

fig, axes = plt.subplots(nrows=1, ncols=2,figsize=(10, 4)) #sets fig and axes variables
women = df[df['Sex']=='female'] #sets women variable to data in women column
men = df[df['Sex']=='male'] #sets men variable to data in men column
ax = sns.histplot(women[women['Survived']==1].Age.dropna(), bins=18, label = 'survived', ax = axes[0], kde =False) #plots women survived graph
ax = sns.histplot(women[women['Survived']==0].Age.dropna(), bins=40, label = 'not survived', ax = axes[0], kde =False) #plots women not survived graph
ax.legend() #creates legend
ax.set_title('Female') #sets title of graph
ax = sns.histplot(men[men['Survived']==1].Age.dropna(), bins=18, label = 'survived', ax = axes[1], kde = False) #plots men survived graph
ax = sns.histplot(men[men['Survived']==0].Age.dropna(), bins=40, label = 'not survived', ax = axes[1], kde = False) #plots men not survived graph
ax.legend() #creates legend
_ = ax.set_title('Male') #sets title of graph
plt.show() #Displays the visualizations


import sys
import json
import yaml

def transformed(data, kind):
    if kind == "secrets":
        if "metadata" in data and "annotations" in data["metadata"] and "kubernetes.io/service-account.name" in data["metadata"]["annotations"]: return None

    if kind == "configmaps":
        if data["metadata"]["name"] in ["config-service-cabundle", "config-trusted-cabundle", "kube-root-ca.crt", "openshift-service-ca.crt"]: return None

    if kind == "serviceaccounts":
        if data["metadata"]["name"] in ["default", "builder", "deployer"]: return None

    if "metadata" in data and "ownerReferences" in data["metadata"]: return None

    if kind in ["rolebindings.rbac.authorization.k8s.io","rolebindings.authorization.openshift.io"]:
        if data["metadata"]["name"] in ["system:deployers", "system:image-builders", "system:image-pullers"]: return None
        if "subjects" in data and len(data["subjects"]) == 1 and "name" in data["subjects"][0] and data["subjects"][0]["name"] == "kube:admin": return None

    if kind == "endpoints": return None

    if "metadata" in data:
        if "creationTimestamp" in data["metadata"]: del data["metadata"]["creationTimestamp"]
        if "managedFields" in data["metadata"]: del data["metadata"]["managedFields"]
        if "resourceVersion" in data["metadata"]: del data["metadata"]["resourceVersion"]
        if "uid" in data["metadata"]: del data["metadata"]["uid"]

    if "status" in data: del data["status"]

    return data

if __name__ == "__main__":
    name = sys.argv[1]
    kind = sys.argv[2]
    path = sys.argv[3]
    f = open(path+name+".json")
    data = json.load(f)
    f.close()
    transformed = transformed(data, kind)
    if transformed is not None:
        file = open(path+name+".yaml", 'w')
        yaml.dump(transformed, file, default_flow_style=False)
        file.close()


import argparse
import os
import numpy as np
import math
import sys

import torchvision.transforms as transforms
from torchvision.utils import save_image

from torch.utils.data import DataLoader
from torchvision import datasets as  datasets
from torch.autograd import Variable

import torch.nn as nn
import torch.nn.functional as F
import torch.autograd as autograd
import torch

from torchvision.utils import make_grid
import matplotlib.pyplot as plt

lastCheckpoint = 514

lr = 0.0002
b1 = 0.5
b2  = 0.999
latent_dim = 100

img_height = 512
img_width  = 272
channels = 3
img_shape = (channels, img_height,  img_width )

outputDir = "Results"
CheckpointsDir = "state_dict"

os.makedirs("{}".format(output), exist_ok=True)

class Generator(nn.Module):
  def __init__(self, input_dim=100, channels=3, memory=8):
    super().__init__()
    self.input_dim = input_dim
    self.decoder = nn.Sequential(  # fully convolutional model

      # input is Z, going into a convolution
      nn.ConvTranspose2d( input_dim, memory * 64, 4, 1, 0, bias=False),
      nn.BatchNorm2d(memory * 64),
      nn.ReLU(True),
      # state size.  4 x 4
      nn.ConvTranspose2d(memory * 64, memory * 32, 4, 2, 1, bias=False),
      nn.BatchNorm2d(memory * 32),
      nn.ReLU(True),
      # state size. 8 x 8
      nn.ConvTranspose2d( memory * 32, memory * 16, 4, 2, 1, bias=False),
      nn.BatchNorm2d(memory * 16),
      nn.ReLU(True),
      # state size. 16 x 16
      nn.ConvTranspose2d( memory * 16, memory * 8, 4, (2,1), 1, bias=False),
      nn.BatchNorm2d(memory * 8),
      nn.ReLU(True),
      # state size. 32 x 64
      nn.ConvTranspose2d( memory * 8, memory * 4, 4, 2, 1, bias=False),
      nn.BatchNorm2d(memory * 4),
      nn.ReLU(True),
      # state size. 64 x 128
      nn.ConvTranspose2d( memory * 4, memory * 2, 4, 2, 1, bias=False),
      nn.BatchNorm2d(memory * 2),
      nn.ReLU(True),
      # state size. 128 x 256
      nn.ConvTranspose2d( memory * 2, memory, 4, 2, 1, bias=False),
      nn.BatchNorm2d(memory),
      nn.ReLU(True),
      # state size.  256 x 512
      nn.ConvTranspose2d( memory, channels, 4, 2, 1, bias=False),
      nn.Tanh()
      # state size. 512 x 256
    )

  def forward(self, latent_code):
    return self.decoder(latent_code.view(-1, self.input_dim, 1, 1))

class Discriminator(nn.Module):
  def __init__(self, channels=3, memory=8):
    super().__init__()
    self.features = nn.Sequential(  # fully convolutional model

      nn.Conv2d(channels, memory, 4, (2,1), 1, bias=False),
      nn.LeakyReLU(0.2, inplace=True),

      nn.Conv2d(memory, memory * 2, 4, 2, 1, bias=False),
      nn.BatchNorm2d(memory * 2),
      nn.LeakyReLU(0.2, inplace=True),

      nn.Conv2d(memory * 2, memory * 4, 4, 2, 1, bias=False),
      nn.BatchNorm2d(memory * 4),
      nn.LeakyReLU(0.2, inplace=True),

      nn.Conv2d(memory * 4, memory * 8, 4, 2, 1, bias=False),
      nn.BatchNorm2d(memory * 8),
      nn.LeakyReLU(0.2, inplace=True),

      nn.Conv2d(memory * 8, memory * 16, 4, 2, 1, bias=False),
      nn.BatchNorm2d(memory * 16),
      nn.LeakyReLU(0.2, inplace=True),

      nn.Conv2d(memory * 16, memory * 32, 4, 2, 1, bias=False),
      nn.BatchNorm2d(memory * 32),
      nn.LeakyReLU(0.2, inplace=True),

      nn.Conv2d(memory * 32, memory * 64, 4, 2, 1, bias=False),
      nn.BatchNorm2d(memory * 64),
      nn.LeakyReLU(0.2, inplace=True),

      nn.Conv2d(memory * 64, 1, 4, 1, 0, bias=False),

    )

  def forward(self, images):

    return self.features(images).flatten(1).mean(1, keepdim=True)

# Loss weight for gradient penalty
lambda_gp = 10

if torch.cuda.device_count() > 0:
    device = torch.device('cuda:0')
else:
    device = torch.device('cpu')

# Initialize generator and discriminator
generator = Generator().to(device)
discriminator = Discriminator().to(device)

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

def load_model(epoch):
    state = torch.load(f'state_dict/state_dict_{epoch}.pt')
    generator.load_state_dict(state['generator'])
    discriminator.load_state_dict(state['discriminator'])
    #optimizer_G.load_state_dict(state['g_optimizer'])
    #optimizer_D.load_state_dict(state['d_optimizer'])

load_model(lastCheckpoint)

def plot_image_grid(images, columns=8, ax=None, show=True):
  if ax is None:
    _, ax = plt.subplots(figsize=(10, 10))
  image_grid = make_grid(images.detach().cpu(), columns, normalize=True)
  ax.imshow(image_grid.permute(1, 2, 0), interpolation='nearest')
  ax.axis('off')
  if show:
    plt.show(ax.figure)
  return ax

# every time you run this, you will plot new randomly generated samples
with torch.no_grad():
  generator.train(False)
  batch_size = 1
  z = torch.randn(batch_size, 100, device=device)
  plot_image_grid(generator(z))


###############################################################################
#    Programming Project #7
#
#    Algorithm
#        display header
#        call function to open file
#            prompt for strings (a file name)
#            loop while the input is incorrect
#                reprompt for strings (a file name)
#        call function to open file
#            prompt for strings (a file name)
#            loop while the input is incorrect
#                reprompt for strings (a file name)
#        display choices
#        prompt for an integer (a choice number)
#        loop while the input is incorrect
#            reprompt for an integer (a choice number)
#        call a fuction chosen corresponding to the choice no.
#            prompt for input(s) depending on the function called
#            loop while the input(s) is/are incorrect
#                reprompt for input(s)
#            display output(s)
#        exit the program if the input for choice in lowercase is '4'
###############################################################################

import matplotlib.pyplot as plt #import the module for plotting
import csv                      #import csv module
from operator import itemgetter #import the module for sorting

MIN_YEAR = 2009    #value of minimum year
MAX_YEAR = 2017    #value of maximum year

def open_file_csv(prompt_str):
    '''
    Open a file (.csv).
    No parameters
    Returns: the opened file for reading only (fp).
    '''
    #while loop to reprompt for input if it's incorrect
    while True:

        try:                           #open the file if successful
            fp = open(prompt_str, "r", encoding="utf-8")
            break
        except FileNotFoundError:      #display error message if not successful
            print("File not found! Try Again!")
            prompt_str = input("Enter the travel data file: ")
    return fp

def open_file_txt(prompt_str):
    '''
    Open a file (.txt).
    No parameters
    Returns: the opened file for reading only (fp).
    '''
    #while loop to reprompt for input if it's incorrect
    while True:

        try:                           #open the file if successful
            fp = open(prompt_str, "r", encoding="utf-8")
            break
        except FileNotFoundError:      #display error message if not successful
            print("File not found! Try Again!")
            prompt_str = input("Enter the country code file: ")
    return fp

def read_country_code_file(fp):
    '''
    This function read through fp the returning value of open_file_txt(prompt_str)
    function and get country code and country name in tuples and add them to a list
    and get it sorted by country code.
    fp: returning value of open_file_txt(prompt_str) function.
    Returns: a list containing tuples of country code and country name

    '''
    fp.seek(0)       #start reading file at the beginning
    fp.readline()    #skip the first line
    list_country_code = []  #start list containing tuples of code and country with an empty list
    for line in fp:
        line = line.strip("\n")
        line = line.split("/")
        line = tuple(line)
        list_country_code.append(line)
    return sorted(list_country_code, key=itemgetter(0)) #sort the list by code alphabetically

def read_travel_file(fp):
    '''
    This function read through a csv file based on fp returning value of
    open_file_csv(prompt_str) function and collect data. Then put the wanted data
    into separately lists with different years.
    fp: returning value of open_file_txt(prompt_str) function.
    Returns: a big list containing lists of different years.

    '''
    reader = csv.reader(fp) #read through flie as a csv file
    next(reader, None) #skip reading the first line

    #get data for each year separately in different lists
    list_2009 = []
    list_2010 = []
    list_2011 = []
    list_2012 = []
    list_2013 = []
    list_2014 = []
    list_2015 = []
    list_2016 = []
    list_2017 = []

    travel_file_list = [] #create travel_file_list of year lists

    #so sclicing to get wanted information from the file
    for line in reader:
        year = int(line[0])
        country_name = line[1][:20]
        country_code = line[2]
        num_departures = int(line[3])
        num_arrivals = int(line[4])
        expenditures = float(line[5])
        receipts = float(line[6])

        #convert values of wanted info to the required forms preparing for further calculations
        num_departures = num_departures / 1000
        expenditures = expenditures / 1000000
        num_arrivals = num_arrivals / 1000
        receipts = receipts / 1000000

        #conditions for average expenditures calculation

        if num_departures == 0:
            avg_expenditures = 0
        else:

            avg_expenditures = round((expenditures/num_departures)*1000, 2)

        #conditions for average arrivals calculation

        if num_arrivals == 0:
            avg_arrivals = 0

        else:

            avg_arrivals = round((receipts/num_arrivals)*1000, 2)

        #create tuples of wanted information
        wanted_info = year, country_name, country_code, num_arrivals, num_departures, expenditures, receipts, avg_expenditures, avg_arrivals
        tuples_info = tuple(wanted_info)

        #add those tuples into separately list of different years
        if tuples_info[0] == 2009:
            list_2009.append(tuples_info)
        elif tuples_info[0] == 2010:
            list_2010.append(tuples_info)
        elif tuples_info[0] == 2011:
            list_2011.append(tuples_info)
        elif tuples_info[0] == 2012:
            list_2012.append(tuples_info)
        elif tuples_info[0] == 2013:
            list_2013.append(tuples_info)
        elif tuples_info[0] == 2014:
            list_2014.append(tuples_info)
        elif tuples_info[0] == 2015:
            list_2015.append(tuples_info)
        elif tuples_info[0] == 2016:
            list_2016.append(tuples_info)
        elif tuples_info[0] == 2017:
            list_2017.append(tuples_info)

    #sort each list by country name
    travel_file_list.append(sorted(list_2009, key=itemgetter(1)))
    travel_file_list.append(sorted(list_2010, key=itemgetter(1)))
    travel_file_list.append(sorted(list_2011, key=itemgetter(1)))
    travel_file_list.append(sorted(list_2012, key=itemgetter(1)))
    travel_file_list.append(sorted(list_2013, key=itemgetter(1)))
    travel_file_list.append(sorted(list_2014, key=itemgetter(1)))
    travel_file_list.append(sorted(list_2015, key=itemgetter(1)))
    travel_file_list.append(sorted(list_2016, key=itemgetter(1)))
    travel_file_list.append(sorted(list_2017, key=itemgetter(1)))

    return travel_file_list

def get_country_code_data(country_code, data_list):
    '''
    This function create a list of data for a country based on the input of
    country_code.
    country_code: three upper case letters refering to a country name (str)
    data_list: returning value of read_travel_file(fp) function
    Returns: list of data for a country or None

    '''
    country_code_list = []
    for lists in data_list:
        for tup in lists:
            if country_code in tup:
                country_code_list.append(tup)
            else:
                continue
    try:
        country_code_list != []
        return country_code_list
    except:
        return None

def display_country_data(country_list):
    '''
    This function displays data for a country.
    country_list: returning value of get_country_code_data(country_code, data_list) function
    Returns: None

    '''
    sum_depart = 0
    sum_arrivals = 0
    sum_ex = 0
    sum_rec = 0
    for tup in country_list:
        sum_depart += tup[4]
        sum_arrivals += tup[3]
        sum_ex += tup[5]
        sum_rec += tup[6]
        print("{:<6d} {:>15,.2f} {:>15,.2f} {:>15,.2f} {:>15,.2f}".format(tup[0],tup[4],tup[3], tup[5],tup[6]))
    print("\n{:<6s} {:>15,.2f} {:>15,.2f} {:>15,.2f} {:>15,.2f}".format("Total",sum_depart,sum_arrivals,sum_ex,sum_rec))

def create_year_list(year_input, data_list):
    '''
    This function create list of yearly data.
    year_input: year (int)
    data_list: returning value of read_travel_file(fp) function
    Returns: a list of data for a year
    '''
    year_list = []
    for lists in data_list:
        for tup in lists:
            if year_input == tup[0]:
                year_list.append(tup)
            else:
                continue
    return year_list

def display_year_data(year_list):
    '''
    This function display data for a year from year_list plus sum for every data
    at the bottom line.
    year_list: returning value of create_year_list(year_input, data_list) function.
    Returns: None
    '''
    sum_depart = 0
    sum_arrivals = 0
    sum_ex = 0
    sum_rec = 0
    for tup in year_list:
        sum_depart += tup[4]
        sum_arrivals += tup[3]
        sum_ex += tup[5]
        sum_rec += tup[6]
        print("{:<20s} {:>15,.2f} {:>15,.2f} {:>15,.2f} {:>15,.2f}".format(tup[1],tup[4],tup[3], tup[5],tup[6]))
    print("\n{:<20s} {:>15,.2f} {:>15,.2f} {:>15,.2f} {:>15,.2f}".format("Total",sum_depart,sum_arrivals,sum_ex,sum_rec))

def prepare_bar_plot(year_list):
    '''
    This function prepares data for bar plotting.
    year_list: returning value of create_year_list(year_input, data_list) function.
    Returns: average expenditures list and average reciepts list sorted by average expenditures
    and average reciepts in descending order.
    '''
    list_avg_ex = []
    count = 0
    year_list = sorted(year_list, key = lambda tup: tup[7], reverse = True)

    for tup in year_list:

            country_name = tup[1]
            avg_ex = tup[7]
            info_avg_ex = country_name, avg_ex
            tuple_avg_ex = tuple(info_avg_ex)
            list_avg_ex.append(tuple_avg_ex)
            count += 1
            if count == 20:
                break

    list_avg_ex = sorted(list_avg_ex, key = itemgetter(1), reverse = True)

    list_avg_rec = []
    count = 0
    year_list = sorted(year_list, key = lambda tup: tup[8], reverse = True)

    for tup in year_list:

            country_name = tup[1]
            avg_rec = tup[8]
            info_avg_rec = country_name, avg_rec
            tuple_avg_rec = tuple(info_avg_rec)
            list_avg_rec.append(tuple_avg_rec)
            count += 1
            if count == 20:
                break

    list_avg_rec = sorted(list_avg_rec, key = itemgetter(1), reverse = True)

    return list_avg_ex, list_avg_rec

def prepare_line_plot(country_list):
    '''
    This function prepares data for line plotting.
    country_list: returning value of get_country_code_data(country_code, data_list) function.
    Returns: average expenditures list and average reciepts list
    '''

    list_avg_ex = []
    list_avg_rec = []
    for tup in country_list:
        avg_ex = tup[7]
        list_avg_ex.append(avg_ex)
        avg_rec = tup[8]
        list_avg_rec.append(avg_rec)

    return list_avg_ex, list_avg_rec

def plot_bar_data(expend_list, receipt_list, year):
    '''
        This function plots the the top 20 countries with the highest average
        expenditures and the top 20 countries with the highest receipts.

        Returns: None

    '''

    # prepare the columns
    countries_expend = [elem[0] for elem in expend_list]
    values_expend = [elem[1] for elem in expend_list]

    countries_receipt = [elem[0] for elem in receipt_list]
    values_receipt = [elem[1] for elem in receipt_list]

    # Average expenditures

    x = range(20) # top 20 countries are to be plotted.

    fig, axs = plt.subplots(2, 1,figsize=(7,10))
    title = "Top 20 countries with highest average expenditures {:4d}".format(year)
    axs[0].set_title(title)
    axs[0].bar(x, values_expend, width=0.4, color='b')
    axs[0].set_ylabel("Avg. Expenditures (US dollar)")
    axs[0].set_xticks(x)
    axs[0].set_xticklabels(countries_expend , rotation='90')

    # Average receipt
    title = "Top 20 countries with highest average receipt  {:4d}".format(year)
    axs[1].set_title(title)
    axs[1].set_ylabel("Avg. Receipts (US dollar)")
    axs[1].bar(x, values_receipt, width=0.4, color='b')
    axs[1].set_xticks(x)
    axs[1].set_xticklabels(countries_receipt , rotation='90')
    fig.tight_layout()
    plt.show()

    ##comment the previous line and uncomment the following two lines when trying to pass Test 4
#    fig.savefig('avg_expense_receipts.png',dpi=100)
#    fig.clf()
def plot_line_data(country_code, expend_list, receipt_list):
    '''
        Plot the line plot for the expenditures and receipts for the
        country between 2009 and 2017

        Returns: None
    '''

    title = "Average expenditures and receipts for {} between 2009 and 2017".format(country_code)
    years = range(MIN_YEAR, MAX_YEAR+1)
    fig, axs = plt.subplots(figsize=(7,5))
    axs.set_title(title)
    axs.set_ylabel("Cost (US dollar)")
    axs.plot(years, expend_list, years, receipt_list)
    axs.legend(['Expenditures','Receipt'])

    plt.show()

    ##comment the previous line and uncomment the following two lines when trying to pass Test 4
#    fig.savefig('line.png',dpi=100)
#    fig.clf()

def main():
    '''
        This function is the main function that is executed automatically
        while the program is running.
        no parameters
        Returns: the result depends on the loop and conditions within
        the function.
    '''
    BANNER = "\nInternational Travel Data Viewer\
\n\nThis program reads and displays departures, arrivals, expenditures, "\
"and receipts for international travels made between 2009 and 2017."
    OPTION = "\nMenu\
\n\t1: Display data by year\
\n\t2: Display data by country\
\n\t3: Display country codes\
\n\t4: Stop the Program\
\n\n\tEnter option number: "

    print(BANNER)

    file_name_travel = input("Enter the travel data file: ")
    fp = open_file_csv(file_name_travel)    #open csv file
    data_list = read_travel_file(fp)

    file_name_cc = input("Enter the country code file: ") #open txt file
    fp = open_file_txt(file_name_cc)
    list_country_code = read_country_code_file(fp)

    choice = input(OPTION)  #input a choice

    #while loop to reprompt for a choice number
    while choice.lower() != '4':    #end the program when input is '6'

        #conditions for choice #1: Top sites by country
        if choice == '1':
            year_input = int(input("Enter year: "))
            li_year = [2009, 2010, 2011, 2012, 2013, 2014, 2015, 2016, 2017]
            while True:
                if year_input not in li_year:
                    print("Year needs to be between 2009 and 2017. Try Again!")
                    year_input = int(input("Enter year: "))
                    continue

                elif year_input in li_year:
                    break

            year_list = create_year_list(year_input, data_list)
            print("\t\t\t\tTravel Data for", year_input)
            print("{:<20s} {:>15s} {:>15s} {:>15s} {:>15s}".format("Country Name","Departures","Arrivals","Expenditures","Receipts"))
            print("{:<20s} {:>15s} {:>15s} {:>15s} {:>15s}".format(" ","(thousands)","(thousands)","(millions)","(millions)"))
            display_year_data(year_list)

            #ask users for plotting
            plot = input("Do you want to plot (yes/no)?").lower()
            if plot == 'yes':
                expend_list, receipt_list = prepare_bar_plot(year_list)
                year = year_input
                plot_bar_data(expend_list, receipt_list, year)
                choice = input(OPTION)
            elif plot == 'no':
                choice = input(OPTION)

        #conditions for choice #2: Display data by country
        elif choice == '2':
            code_list = []

            for lists in data_list:
                for tup in lists:
                    cc = tup[2]
                    code_list.append(cc)

            country_code = input("Enter country code: ").upper()

            if country_code not in code_list:
                print("Country code is not found! Try Again!")
                continue

            elif country_code in code_list:
                country_list = get_country_code_data(country_code, data_list)

            #get country name for header
            country_code_list = read_country_code_file(fp)
            country_name = " "
            for tup in country_code_list:
                if tup[0] == country_code:
                    country_name = tup[1]
                    break

            print("\t\t\t   Travel Data for", country_name)
            print("{:<6s} {:>15s} {:>15s} {:>15s} {:>15s}".format("Year","Departures","Arrivals","Expenditures","Receipts"))
            print("{:<6s} {:>15s} {:>15s} {:>15s} {:>15s}".format(" ","(thousands)","(thousands)","(millions)","(millions)"))

            display_country_data(country_list)

            #ask users for plotting
            plot = input("Do you want to plot (yes/no)?").lower()
            if plot == 'yes':
                prepare_line_plot(country_list)
                expend_list_forline, receipt_list_forline = prepare_line_plot(country_list)
                plot_line_data(country_code, expend_list_forline, receipt_list_forline)

                choice = input(OPTION)

            elif plot == 'no':
                choice = input(OPTION)

        #conditions for choice #3: Display country codes
        elif choice == '3':
            print("\nCountry Code Reference")
            print("{:15s}{:25s}".format("Country Code", "Country Name"))
            for tup in list_country_code:
                print("{:15s}{:25s}".format(tup[0], tup[1]))

            choice = input(OPTION)

        #display error message for invalid input choice number
        #reprompt for a choice number
        else:
            print("Invalid option. Try Again!")
            choice = input(OPTION)

    #quit the program if the lowercase of the choice is '4'
    if choice.lower() == '4':
        print("\nThanks for using this program!")

# this was added to make functions to be tested individually on mirmir
if __name__ == "__main__":
    main()


import json
import logging
import os
import shutil
import socket
import subprocess
import sys
import time

import docker

client = docker.from_env()
api_client = docker.APIClient(base_url="unix://var/run/docker.sock")
hostname = socket.gethostname()
current_container = client.containers.get(hostname)

def run(*cmd):
    subprocess.run(cmd, check=True)

# Populate directory structure in data directory.
nginx_data_path = "/root/https/nginx"
nginx_config_path = os.path.join(nginx_data_path, "conf.d")
nginx_webroot_path = os.path.join(nginx_data_path, "webroot")
certbot_data_path = "/root/https/certbot"
logs_path = "root/https/logs"
run("mkdir", "-p", nginx_data_path)
run("mkdir", "-p", nginx_config_path)
run("mkdir", "-p", nginx_webroot_path)
run("mkdir", "-p", certbot_data_path)
run("mkdir", "-p", logs_path)

# Get logger
log_file_path = os.path.join(logs_path, "logs.txt")
logging.basicConfig(level=logging.INFO, stream=sys.stdout)
logger = logging.getLogger("https")

def container_run(container, cmd):
    (exit_code, output) = container.exec_run(cmd)
    logger.info(output.decode("utf-8").rstrip())
    if exit_code != 0:
        sys.exit(exit_code)

# Load configurations.
config_file_path = "/root/https/config.json"
if not os.path.exists(config_file_path):
    config_file_path = "/root/docker/config.json"
with open(config_file_path) as file:
    config = json.load(file)
data_path = config["data_path"]
network_name = config["network_name"]
container_prefix = config["container_prefix"]
domain_name = config["domain_name"]
email = config["email"]
nginx_image_name = config["nginx_image_name"]
certbot_image_name = config["certbot_image_name"]

# Copy Nginx configurations.
nginx_default_config_path = os.path.join(nginx_config_path, "nginx.conf")
nginx_http_config_path = os.path.join(nginx_config_path, "http.conf")
nginx_https_config_path = os.path.join(nginx_config_path, "https.conf")
with open(nginx_default_config_path, "w") as file:
    file.write("include /root/nginx/conf.d/http.conf;\n")
run("cp", "/root/docker/http.conf", nginx_http_config_path)
if not os.path.exists(nginx_https_config_path):
    run("cp", "/root/docker/https.conf", nginx_https_config_path)

# Pull images
logger.info(f"Pulling {nginx_image_name}:")
for line in api_client.pull(nginx_image_name, stream=True, decode=True):
    logger.info(json.dumps(line, indent=4))
logger.info(f"Pulling {certbot_image_name}:")
for line in api_client.pull(certbot_image_name, stream=True, decode=True):
    logger.info(json.dumps(line, indent=4))

# Remove existing containers.
web_container_name = container_prefix + "web"
nginx_container_name = container_prefix + "nginx"
renew_container_name = container_prefix + "renew"
for container_name in [web_container_name,
                       nginx_container_name,
                       renew_container_name]:
    containers = client.containers.list(all=True,
                                        filters={"name": container_name})
    for container in containers:
        logger.info(f"Removing {container_name}")
        container.remove(v=True, force=True)

# Create network
if not client.networks.list(names=[network_name]):
    client.networks.create(network_name)

# Create HTTP Nginx container.
if not domain_name:
    domain_name_file_path = os.path.join(logs_path, "domain_name.txt")
    if os.path.exists(domain_name_file_path):
        with open(domain_name_file_path, "r") as domain_name_file:
            domain_name = domain_name_file.read().strip()
if not domain_name:
    with open(domain_name_file_path, "w") as file:
        pass
    with open(log_file_path, "w") as file:
        pass
    shutil.copytree("/root/web/root",
                    nginx_webroot_path,
                    dirs_exist_ok=True)
    run("cp", "/root/web/nginx/http.conf", nginx_http_config_path)
    web = api_client.create_container(
        current_container.image.id,
        command=["python", "-u", "/root/web/server/main.py"],
        detach=True,
        host_config=api_client.create_host_config(
            binds={data_path: {"bind": "/root/https/",
                               "mode": "rw"}}),
        name=web_container_name,
        networking_config=api_client.create_networking_config({
            network_name: api_client.create_endpoint_config(aliases=[
                                                            "web"])
        }),
        volumes=["/root/https/"])
    api_client.start(container=web.get('Id'))
logger.info(f"Running {nginx_image_name}.")
host_nginx_config_path = os.path.join(data_path, "nginx", "conf.d")
host_nginx_webroot_path = os.path.join(data_path, "nginx", "webroot")
host_certbot_data_path = os.path.join(data_path, "certbot")
nginx_container = client.containers.run(
    nginx_image_name,
    entrypoint="nginx",
    command=["-c", "/root/nginx/conf.d/nginx.conf",
             "-g", "daemon off;"],
    detach=True,
    links={web_container_name: "web"},
    name=nginx_container_name,
    network=network_name,
    ports={80: 80,
           443: 443},
    restart_policy={"Name": "always",
                    "MaximumRetryCount": 0},
    stream=True,
    volumes={host_nginx_config_path: {"bind": "/root/nginx/conf.d",
                                      "mode": "rw"},
             host_nginx_webroot_path: {"bind": "/var/www",
                                       "mode": "rw"},
             host_certbot_data_path: {"bind": "/etc/letsencrypt",
                                      "mode": "rw"},
             "/root/sockets/": {"bind": "/root/sockets/",
                                "mode": "rw"}})
if not domain_name:
    logger.info("Waiting web request (enter domain name in browser).")
    domain_name_file = open(domain_name_file_path, "r")
    while True:
        line = domain_name_file.readline()
        if line:
            domain_name = line.strip()
            break
        else:
            time.sleep(0.01)

# Run certbot to obtain certificates.
logger.addHandler(logging.FileHandler(log_file_path))
logger.info(f"Running {certbot_image_name}.")
if email:
    email_options = ["--email", email]
else:
    email_options = ["--register-unsafely-without-email"]
certbot_container = client.containers.run(
    certbot_image_name,
    detach=True,
    volumes={host_nginx_webroot_path: {"bind": "/var/www",
                                       "mode": "rw"},
             host_certbot_data_path: {"bind": "/etc/letsencrypt",
                                      "mode": "rw"}},
    command=["certonly",
             "--webroot",
             "--non-interactive",
             "--domain", domain_name,
             "--agree-tos",
             *email_options,
             "--no-eff-email",
             "--webroot-path", "/var/www"])
for line in certbot_container.logs(stream=True):
    logger.info(line.decode("utf-8").rstrip())
status = certbot_container.wait()["StatusCode"]
certbot_container.remove(v=True, force=True)
if status != 0:
    sys.exit(status)

# Create symbolic links to certificate files and reload nginx.
container_run(nginx_container, ["mkdir", "-p", "/root/nginx/ssl"])
for file_name in ["fullchain.pem", "privkey.pem", "chain.pem"]:
    cmd = ["ln", "-sf",
           f"/etc/letsencrypt/live/{domain_name}/{file_name}",
           f"/root/nginx/ssl/{file_name}"]
    container_run(nginx_container, cmd)
with open(nginx_default_config_path, "w") as file:
    file.write("include /root/nginx/conf.d/https.conf;\n")
container_run(nginx_container, ["nginx", "-s", "reload"])

# Run renew.py using the current image.
client.containers.run(
    current_container.image.id,
    command=["python", "-u", "/root/docker/renew.py"],
    detach=True,
    name=renew_container_name,
    restart_policy={"Name": "always",
                    "MaximumRetryCount": 0},
    volumes={data_path: {"bind": "/root/https",
                         "mode": "rw"},
             "/var/run/docker.sock": {"bind": "/var/run/docker.sock",
                                      "mode": "rw"}})

containers = client.containers.list(all=True,
                                    filters={"name": web_container_name})
for container in containers:
    container.remove(v=True, force=True)

logger.info("Done.")


from transformers import ViTModel, PreTrainedModel, PretrainedConfig
from torch import nn
from transformers.activations import ACT2FN
from torch.nn import CrossEntropyLoss, MSELoss
import math
from transformers.modeling_outputs import (
    BaseModelOutput,
    BaseModelOutputWithPooling,
    SequenceClassifierOutput,
)
import torch

class ViViTConfig(PretrainedConfig):
    model_type = "vivit"

    def __init__(
        self,
        hidden_size=768,
        num_hidden_layers=12,
        num_attention_heads=12,
        intermediate_size=3072,
        hidden_act="gelu",
        hidden_dropout_prob=0.0,
        attention_probs_dropout_prob=0.0,
        initializer_range=0.02,
        layer_norm_eps=1e-12,
        is_encoder_decoder=False,
        video_size=(30, 224, 224),
        tubelet_size=(2, 16, 16),
        num_channels=3,
        qkv_bias=True,
        **kwargs,
    ):
        super().__init__(**kwargs)

        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.intermediate_size = intermediate_size
        self.hidden_act = hidden_act
        self.hidden_dropout_prob = hidden_dropout_prob
        self.attention_probs_dropout_prob = attention_probs_dropout_prob
        self.initializer_range = initializer_range
        self.layer_norm_eps = layer_norm_eps

        self.video_size = video_size
        self.tubelet_size = tubelet_size
        self.num_channels = num_channels
        self.qkv_bias = qkv_bias

class TubeletEmbeddings(nn.Module):
    """
    Video to Tubelet Embedding.

    """

    def __init__(self, video_size, patch_size, num_channels=3, embed_dim=768):
        super().__init__()
        self.video_size = video_size
        self.patch_size = patch_size
        self.num_patches = (
            (video_size[2] // patch_size[2])
            * (video_size[1] // patch_size[1])
            * (video_size[0] // patch_size[0])
        )
        self.embed_dim = embed_dim

        self.projection = nn.Conv3d(
            num_channels, embed_dim, kernel_size=patch_size, stride=patch_size
        )

    def forward(self, pixel_values, interpolate_pos_encoding=False):
        batch_size, num_channels, num_frames, height, width = pixel_values.shape
        if not interpolate_pos_encoding:
            if height != self.video_size[-2] or width != self.video_size[-1]:
                raise ValueError(
                    f"Input image size ({height}*{width}) doesn't match model ({self.video_size[-2]}*{self.video_size[-1]})."
                )

        # TODO: check this
        x = self.projection(pixel_values).flatten(2).transpose(1, 2)
        return x

class ViViTEmbeddings(nn.Module):
    def __init__(self, config):
        super().__init__()

        self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
        self.patch_embeddings = TubeletEmbeddings(
            config.video_size,
            config.tubelet_size,
            config.num_channels,
            config.hidden_size,
        )

        self.position_embeddings = nn.Parameter(
            torch.zeros(1, self.patch_embeddings.num_patches + 1, config.hidden_size)
        )
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.config = config

    def interpolate_pos_encoding(self, embeddings, height, width):
        # TODO: Implement interpolation
        raise NotImplementedError()

    def forward(self, pixel_values, interpolate_pos_encoding=False):
        batch_size, num_frames, num_channels, height, width = pixel_values.shape
        embeddings = self.patch_embeddings(
            pixel_values, interpolate_pos_encoding=interpolate_pos_encoding
        )

        # add the [CLS] token to the embedded patch tokens
        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
        embeddings = torch.cat((cls_tokens, embeddings), dim=1)

        # add positional encoding to each token
        if interpolate_pos_encoding:
            embeddings = embeddings + self.interpolate_pos_encoding(
                embeddings, height, width
            )
        else:
            embeddings = embeddings + self.position_embeddings

        embeddings = self.dropout(embeddings)

        return embeddings

class ViViTSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(
            config, "embedding_size"
        ):
            raise ValueError(
                f"The hidden size {config.hidden_size,} is not a multiple of the number of attention "
                f"heads {config.num_attention_heads}."
            )

        self.num_attention_heads = config.num_attention_heads
        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
        self.all_head_size = self.num_attention_heads * self.attention_head_size

        self.query = nn.Linear(
            config.hidden_size, self.all_head_size, bias=config.qkv_bias
        )
        self.key = nn.Linear(
            config.hidden_size, self.all_head_size, bias=config.qkv_bias
        )
        self.value = nn.Linear(
            config.hidden_size, self.all_head_size, bias=config.qkv_bias
        )

        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)

    def transpose_for_scores(self, x):
        new_x_shape = x.size()[:-1] + (
            self.num_attention_heads,
            self.attention_head_size,
        )
        x = x.view(*new_x_shape)
        return x.permute(0, 2, 1, 3)

    def forward(self, hidden_states, head_mask=None, output_attentions=False):
        mixed_query_layer = self.query(hidden_states)

        key_layer = self.transpose_for_scores(self.key(hidden_states))
        value_layer = self.transpose_for_scores(self.value(hidden_states))
        query_layer = self.transpose_for_scores(mixed_query_layer)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))

        attention_scores = attention_scores / math.sqrt(self.attention_head_size)

        # Normalize the attention scores to probabilities.
        attention_probs = nn.functional.softmax(attention_scores, dim=-1)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        context_layer = torch.matmul(attention_probs, value_layer)

        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(*new_context_layer_shape)

        outputs = (
            (context_layer, attention_probs) if output_attentions else (context_layer,)
        )

        return outputs

class ViViTSelfOutput(nn.Module):
    """
    The residual connection is defined in ViTLayer instead of here (as is the case with other models), due to the
    layernorm applied before each block.
    """

    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states, input_tensor):

        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)

        return hidden_states

class ViViTAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.attention = ViViTSelfAttention(config)
        self.output = ViViTSelfOutput(config)
        self.pruned_heads = set()

    def prune_heads(self, heads):
        if len(heads) == 0:
            return
        heads, index = find_pruneable_heads_and_indices(
            heads,
            self.attention.num_attention_heads,
            self.attention.attention_head_size,
            self.pruned_heads,
        )

        # Prune linear layers
        self.attention.query = prune_linear_layer(self.attention.query, index)
        self.attention.key = prune_linear_layer(self.attention.key, index)
        self.attention.value = prune_linear_layer(self.attention.value, index)
        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)

        # Update hyper params and store pruned heads
        self.attention.num_attention_heads = self.attention.num_attention_heads - len(
            heads
        )
        self.attention.all_head_size = (
            self.attention.attention_head_size * self.attention.num_attention_heads
        )
        self.pruned_heads = self.pruned_heads.union(heads)

    def forward(self, hidden_states, head_mask=None, output_attentions=False):
        self_outputs = self.attention(hidden_states, head_mask, output_attentions)

        attention_output = self.output(self_outputs[0], hidden_states)

        outputs = (attention_output,) + self_outputs[
            1:
        ]  # add attentions if we output them
        return outputs

class ViViTIntermediate(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
        if isinstance(config.hidden_act, str):
            self.intermediate_act_fn = ACT2FN[config.hidden_act]
        else:
            self.intermediate_act_fn = config.hidden_act

    def forward(self, hidden_states):

        hidden_states = self.dense(hidden_states)
        hidden_states = self.intermediate_act_fn(hidden_states)

        return hidden_states

class ViViTOutput(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states, input_tensor):
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)

        hidden_states = hidden_states + input_tensor

        return hidden_states

class ViViTLayer(nn.Module):
    """This corresponds to the Block class in the timm implementation."""

    def __init__(self, config):
        super().__init__()
        self.chunk_size_feed_forward = config.chunk_size_feed_forward
        self.seq_len_dim = 1
        self.attention = ViViTAttention(config)
        self.intermediate = ViViTIntermediate(config)
        self.output = ViViTOutput(config)
        self.layernorm_before = nn.LayerNorm(
            config.hidden_size, eps=config.layer_norm_eps
        )
        self.layernorm_after = nn.LayerNorm(
            config.hidden_size, eps=config.layer_norm_eps
        )

    def forward(self, hidden_states, head_mask=None, output_attentions=False):
        self_attention_outputs = self.attention(
            self.layernorm_before(
                hidden_states
            ),  # in ViT, layernorm is applied before self-attention
            head_mask,
            output_attentions=output_attentions,
        )
        attention_output = self_attention_outputs[0]
        outputs = self_attention_outputs[
            1:
        ]  # add self attentions if we output attention weights

        # first residual connection
        hidden_states = attention_output + hidden_states

        # in ViT, layernorm is also applied after self-attention
        layer_output = self.layernorm_after(hidden_states)
        layer_output = self.intermediate(layer_output)

        # second residual connection is done here
        layer_output = self.output(layer_output, hidden_states)

        outputs = (layer_output,) + outputs

        return outputs

class ViViTEncoder(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.layer = nn.ModuleList(
            [ViViTLayer(config) for _ in range(config.num_hidden_layers)]
        )
        self.gradient_checkpointing = False

    def forward(
        self,
        hidden_states,
        head_mask=None,
        output_attentions=False,
        output_hidden_states=False,
        return_dict=True,
    ):
        all_hidden_states = () if output_hidden_states else None
        all_self_attentions = () if output_attentions else None

        for i, layer_module in enumerate(self.layer):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            layer_head_mask = head_mask[i] if head_mask is not None else None

            if self.gradient_checkpointing and self.training:

                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        return module(*inputs, output_attentions)

                    return custom_forward

                layer_outputs = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(layer_module),
                    hidden_states,
                    layer_head_mask,
                )
            else:
                layer_outputs = layer_module(
                    hidden_states, layer_head_mask, output_attentions
                )

            hidden_states = layer_outputs[0]

            if output_attentions:
                all_self_attentions = all_self_attentions + (layer_outputs[1],)

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        if not return_dict:
            return tuple(
                v
                for v in [hidden_states, all_hidden_states, all_self_attentions]
                if v is not None
            )
        return BaseModelOutput(
            last_hidden_state=hidden_states,
            hidden_states=all_hidden_states,
            attentions=all_self_attentions,
        )

class ViViTPreTrainedModel(PreTrainedModel):
    """
    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
    models.
    """

    config_class = ViViTConfig
    base_model_prefix = "vivit"
    main_input_name = "pixel_values"
    supports_gradient_checkpointing = True

    def _init_weights(self, module):
        """Initialize the weights"""
        if isinstance(module, (nn.Linear, nn.Conv3d)):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)

    def _set_gradient_checkpointing(self, module, value=False):
        if isinstance(module, ViViTEncoder):
            module.gradient_checkpointing = value

class ViViTModel(ViViTPreTrainedModel):
    def __init__(self, config, add_pooling_layer=True):
        super().__init__(config)
        self.config = config

        self.embeddings = ViViTEmbeddings(config)
        self.encoder = ViViTEncoder(config)

        self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.pooler = ViViTPooler(config) if add_pooling_layer else None

        # Initialize weights and apply final processing
        self.post_init()

    def get_input_embeddings(self):
        return self.embeddings.patch_embeddings

    def _prune_heads(self, heads_to_prune):
        """
        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
        class PreTrainedModel
        """
        for layer, heads in heads_to_prune.items():
            self.encoder.layer[layer].attention.prune_heads(heads)

    def forward(
        self,
        pixel_values=None,
        head_mask=None,
        output_attentions=None,
        output_hidden_states=None,
        interpolate_pos_encoding=None,
        return_dict=None,
    ):
        output_attentions = (
            output_attentions
            if output_attentions is not None
            else self.config.output_attentions
        )
        output_hidden_states = (
            output_hidden_states
            if output_hidden_states is not None
            else self.config.output_hidden_states
        )
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        if pixel_values is None:
            raise ValueError("You have to specify pixel_values")

        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)

        embedding_output = self.embeddings(
            pixel_values, interpolate_pos_encoding=interpolate_pos_encoding
        )

        encoder_outputs = self.encoder(
            embedding_output,
            head_mask=head_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        sequence_output = encoder_outputs[0]
        sequence_output = self.layernorm(sequence_output)
        pooled_output = (
            self.pooler(sequence_output) if self.pooler is not None else None
        )

        if not return_dict:
            return (sequence_output, pooled_output) + encoder_outputs[1:]

        return BaseModelOutputWithPooling(
            last_hidden_state=sequence_output,
            pooler_output=pooled_output,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
        )

class ViViTForImageClassification(ViViTPreTrainedModel):
    def __init__(self, config):
        super(ViViTForImageClassification, self).__init__(config)

        self.num_labels = config.num_labels
        self.vivit = ViViTModel(config, add_pooling_layer=False)

        # Classifier head
        self.classifier = (
            nn.Linear(config.hidden_size, config.num_labels)
            if config.num_labels > 0
            else nn.Identity()
        )

        # Initialize weights and apply final processing
        self.post_init()

    def forward(
        self,
        pixel_values=None,
        head_mask=None,
        labels=None,
        output_attentions=None,
        output_hidden_states=None,
        interpolate_pos_encoding=None,
        return_dict=None,
    ):
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        outputs = self.vivit(
            pixel_values,
            head_mask=head_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            interpolate_pos_encoding=interpolate_pos_encoding,
            return_dict=return_dict,
        )

        sequence_output = outputs[0]

        logits = self.classifier(sequence_output[:, 0, :])

        loss = None
        if labels is not None:
            if self.num_labels == 1:
                #  We are doing regression
                loss_fct = MSELoss()
                loss = loss_fct(logits.view(-1), labels.view(-1))
            else:
                loss_fct = CrossEntropyLoss()
                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))

        if not return_dict:
            output = (logits,) + outputs[2:]
            return ((loss,) + output) if loss is not None else output

        return SequenceClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

class ViViTPooler(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.activation = nn.Tanh()

    def forward(self, hidden_states):
        # We "pool" the model by simply taking the hidden state corresponding
        # to the first token.
        first_token_tensor = hidden_states[:, 0]
        pooled_output = self.dense(first_token_tensor)
        pooled_output = self.activation(pooled_output)
        return pooled_output