from __future__ import division

import sys
import random
import datetime
import torch
import numpy as np
import torch.nn.functional as F
from torch.autograd import Variable

from . import update
from . import network
from .config import *
from . import OUnoise

sys.path.append("..")
from etc import ALGO_PATH
from . import buffer

class Trainer:
    '''训练器'''

    def __init__(self, state_dim, action_dim, dev, epoch=0, model_save_path=None):
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.ram = buffer.MemoryBuffer(MAX_BUFFER)
        self.noise = OUnoise.OrnsteinUhlenbeckActionNoise(self.action_dim)
        self.device = dev

        # 策略网络（在线和目标）及其优化器
        self.actor = network.Actor(self.state_dim, self.action_dim).to(self.device)
        self.target_actor = network.Actor(self.state_dim, self.action_dim).to(self.device)
        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), LEARNING_RATE)

        # 价值网络（在线和目标）及其优化器
        self.critic = network.Critic(self.state_dim, self.action_dim).to(self.device)
        self.target_critic = network.Critic(self.state_dim, self.action_dim).to(self.device)
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), LEARNING_RATE)

        # 如果当前轮数大于0，则表示以前训练过，则载入训练过的模型
        if epoch > 0:
            print("load models")
            self.load_models(epoch, model_save_path)

        # 用在线网络更新目标网络
        update.hard_update(self.target_actor, self.actor)
        update.hard_update(self.target_critic, self.critic)

    def get_exploitation_action(self, state):
        '''利用模型产生动作'''
        state = Variable(torch.from_numpy(state).to(self.device))
        action = self.target_actor.forward(state).detach()
        return action.data.cpu().numpy()

    def get_exploration_action(self, state):
        '''随机探索产生动作'''
        state = Variable(torch.from_numpy(state).to(self.device))
        action = self.actor.forward(state).detach()
        new_action = action.data.cpu().numpy() + self.noise.sample()
        return new_action

    # 用深度强化学习算法得到智能体动作
    def get_action_from_drl_algo(self, x, y):
        state = [x, y]
        state = np.float32(state)  # 强制转化为浮点类型

        # epsilon的动作随机生成，其他的动作由模型生成
        number = random.randint(0, 1)
        if number < EPSILON:
            action = self.get_exploration_action(state)  # 通过模型来生成动作，利用
        else:
            action = self.get_exploitation_action(state)  # 随机生成动作
        with open("%s/output/action.csv" % ALGO_PATH, "a+") as f:
            # f.write("action from mfddpg is [%f,%f]" % (action[0], action[1]))
            f.write("%f,%f\n" % (action[0], action[1]))
        # print("action from mfddpg is ", action)
        return action

    # 根据动作执行结果，训练一次智能体
    def train_(self, train_step, state_last, action_last, reward_now, state_new, cur_step, mean_action):
        # 添加最新经验，并优化训练一把，再做决策
        self.ram.add(state_last, action_last, reward_now, state_new, mean_action)
        self.optimize(cur_step)
        # print("step in mfddpg is ", cur_step)
        train_step += 1

    def optimize(self, step):
        '''优化'''
        s1, a1, r1, s2, a_ = self.ram.sample(BATCH_SIZE)
        s1 = Variable(torch.from_numpy(s1).to(self.device))
        a1 = Variable(torch.from_numpy(a1).to(self.device))
        r1 = Variable(torch.from_numpy(r1).to(self.device))
        s2 = Variable(torch.from_numpy(s2).to(self.device))
        a_ = Variable(torch.from_numpy(a_).to(self.device))

        a2 = self.target_actor.forward(s2).detach()  # 根据S2得出的下一个动作a2
        # print("a2 is ", a2.size())
        next_val = torch.squeeze(self.target_critic.forward(s2, a2, a_).detach())  # 根据s2和a2得出的目标值
        # print("next_val is ", next_val.size())
        y_expectd = r1 + GAMMA * next_val  # 目标回报值
        # print("y_expected is ", y_expectd.size())

        y_predicted = torch.squeeze(self.critic.forward(s1, a1, a_))  # 主价值网络得到的回报值
        # print("y_predicted is ", y_predicted.size())
        loss_critic = F.smooth_l1_loss(y_predicted, y_expectd)  # 价值损失（回报值的损失）
        self.critic_optimizer.zero_grad()  # 优化器参数置为零
        with open("%s/output/loss_critic.csv" % ALGO_PATH, "a+") as f:
            f.write("%s,%s\n" % (step, loss_critic.item()))  # 将损失值保存到文件中
        loss_critic.backward()  # 反向传播
        self.critic_optimizer.step()  # 修改价值网络参数

        pred_a1 = self.actor.forward(s1)
        loss_actor = -1 * torch.sum(self.critic.forward(s1, pred_a1, a_))  # 决策损失（回报值的损失）
        self.actor_optimizer.zero_grad()
        with open("%s/output/loss_actor.csv" % ALGO_PATH, "a+") as f:
            f.write("%s,%s\n" % (step, loss_actor.item()))  # 将损失值保存到文件中
        loss_actor.backward()  # 反向传播
        self.actor_optimizer.step()  # 修改策略网络参数

        update.soft_update(self.target_actor, self.actor, TAU)
        update.soft_update(self.target_critic, self.critic, TAU)

    def get_models_path(self):
        return "./Models/"

    def save_model(self, episode_count, model_save_path=None):
        '''保存模型'''
        # torch.save(self.target_actor.state_dict(), './Models/' + str(episode_count) + '_actor.pt')
        if model_save_path == None:
            model_save_path = self.get_models_path()
        torch.save(self.target_actor.state_dict(), model_save_path + str(episode_count) + '_actor.pt')
        torch.save(self.target_critic.state_dict(), model_save_path + str(episode_count) + '_critic.pt')
        print("%s：%s Models saved successfully" % (datetime.datetime.now(), episode_count))
        # print("%s：轮数:%s 决策步数:%s  Reward:%.2f" %  (datetime.now(), episode_count, step, reward_now))

    def load_models(self, episode, model_save_path=None):
        '''载入以前训练过的模型, 包括策略网络和价值网络'''
        if model_save_path == None:
            model_save_path = self.get_models_path()
        # self.critic.load_state_dict(torch.load(self.get_models_path() + str(episode) + '_critic.pt'))
        self.critic.load_state_dict(torch.load(model_save_path + str(episode) + '_critic.pt'))
        # self.actor.load_state_dict(torch.load(self.get_models_path() + str(episode) + '_actor.pt'))
        self.actor.load_state_dict(torch.load(model_save_path + str(episode) + '_actor.pt'))
        update.hard_update(self.target_actor, self.actor)
        update.hard_update(self.target_critic, self.critic)
        print("Models loaded successfully")


import os

import numpy as np

import torch
from torch import nn
from boxlist import cat_boxlist, boxlist_iou

import math

INF = 100000000

class SigmoidFocalLoss(nn.Module):
    def __init__(self, gamma, alpha):
        super().__init__()

        self.gamma = gamma
        self.alpha = alpha
        self.eps = 1e-6

    def forward(self, out, target):
        n_class = out.shape[1]
        class_ids = torch.arange(
            1, n_class + 1, dtype=target.dtype, device=target.device
        ).unsqueeze(0)

        t = target.unsqueeze(1)
        p = torch.sigmoid(out)
        p = torch.clamp(p, self.eps, 1-self.eps) # for numerical stability

        gamma = self.gamma
        alpha = self.alpha

        term1 = (1 - p) ** gamma * torch.log(p)
        term2 = p ** gamma * torch.log(1 - p)

        # print(term1.sum(), term2.sum())

        loss = (
            -(t == class_ids).float() * alpha * term1
            - ((t != class_ids) * (t >= 0)).float() * (1 - alpha) * term2
        )

        return loss.sum()

def get_num_gpus():
    return int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1

def reduce_sum(tensor):
    if get_num_gpus() <= 1:
        return tensor
    import torch.distributed as dist
    tensor = tensor.clone()
    dist.all_reduce(tensor, op=dist.reduce_op.SUM)
    return tensor

def cat(tensors, dim=0):
    """
    Efficient version of torch.cat that avoids a copy if there is only a single element in a list
    """
    assert isinstance(tensors, (list, tuple))
    if len(tensors) == 1:
        return tensors[0]
    return torch.cat(tensors, dim)

def permute_and_flatten(layer, N, A, C, H, W):
    layer = layer.view(N, -1, C, H, W)
    layer = layer.permute(0, 3, 4, 1, 2)
    layer = layer.reshape(N, -1, C)
    return layer

def concat_box_prediction_layers(pred_cls, pred_reg):
    pred_cls_flattened = []
    pred_reg_flattened = []
    # for each feature level, permute the outputs to make them be in the
    # same format as the labels. Note that the labels are computed for
    # all feature levels concatenated, so we keep the same representation
    # for the objectness and the pred_reg
    for pred_cls_per_level, pred_reg_per_level in zip(pred_cls, pred_reg):
        N, AxC, H, W = pred_cls_per_level.shape
        Cx16 = pred_reg_per_level.shape[1]
        C = Cx16 // 16
        A = 1
        pred_cls_per_level = permute_and_flatten(
            pred_cls_per_level, N, A, C, H, W
        )
        pred_cls_flattened.append(pred_cls_per_level)

        pred_reg_per_level = permute_and_flatten(
            pred_reg_per_level, N, A, (C*16), H, W
        )
        pred_reg_flattened.append(pred_reg_per_level)
    # concatenate on the first dimension (representing the feature levels), to
    # take into account the way the labels were generated (with all feature maps
    # being concatenated as well)
    pred_cls = cat(pred_cls_flattened, dim=1).reshape(-1, C)
    pred_reg = cat(pred_reg_flattened, dim=1).reshape(-1, C*16)
    return pred_cls, pred_reg

class PoseLoss(object):
    def __init__(self, gamma, alpha, anchor_sizes, anchor_strides, positive_num, positive_lambda,
                    loss_weight_cls, loss_weight_reg, internal_K, diameters, target_coder):
        self.cls_loss_func = SigmoidFocalLoss(gamma, alpha)
        # self.centerness_loss_func = nn.BCEWithLogitsLoss(reduction="sum")
        # self.matcher = Matcher(fg_iou_threshold, bg_iou_threshold, True)
        self.anchor_sizes = anchor_sizes
        self.anchor_strides = anchor_strides
        self.positive_num = positive_num
        self.positive_lambda = positive_lambda
        self.loss_weight_cls = loss_weight_cls
        self.loss_weight_reg = loss_weight_reg
        self.internal_K = internal_K
        self.target_coder = target_coder
        self.diameters = diameters

    def ObjectSpaceLoss(self, pred, target_3D_in_camera_frame, cls_labels, anchors, weight=None):
        if not isinstance(self.diameters, torch.Tensor):
            self.diameters = torch.FloatTensor(self.diameters).to(device=pred.device).view(-1)

        diameter_ext = self.diameters[cls_labels.view(-1,1).repeat(1, 8*3).view(-1, 3, 1)]

        cellNum = pred.shape[0]
        pred_filtered = pred.view(cellNum, -1, 16)[torch.arange(cellNum), cls_labels]

        pred_xy = self.target_coder.decode(pred_filtered, anchors)
        # target_xy = self.target_coder.decode(target, anchors)

        pred_xy = pred_xy.view(-1,2,8).transpose(1,2).contiguous().view(-1,2)

        # construct normalized 2d
        B = torch.inverse(self.internal_K).mm(torch.cat((pred_xy.t(), torch.ones_like(pred_xy[:,0]).view(1,-1)), dim=0)).t()
        # compute projection matrices
        P = torch.bmm(B.view(-1, 3, 1), B.view(-1, 1, 3)) / torch.bmm(B.view(-1, 1, 3), B.view(-1, 3, 1))

        target_3D_in_camera_frame = target_3D_in_camera_frame.view(-1, 3, 1)
        px = torch.bmm(P, target_3D_in_camera_frame)

        target_3D_in_camera_frame = target_3D_in_camera_frame / diameter_ext
        px = px / diameter_ext
        scaling_factor = 50 # 0.02d
        losses = nn.SmoothL1Loss(reduction='none')(scaling_factor * px, scaling_factor * target_3D_in_camera_frame).view(cellNum, -1).mean(dim=1)
        losses = losses / scaling_factor

        if weight is not None and weight.sum() > 0:
            return (losses * weight).sum()
        else:
            assert losses.numel() != 0
            return losses.sum()

    def prepare_targets(self, targets, anchors):
        cls_labels = []
        reg_targets = []
        aux_raw_boxes = []
        aux_3D_in_camera_frame = []
        level_cnt = len(anchors[0])
        for im_i in range(len(targets)):
            pose_targets_per_im = targets[im_i]
            bbox_targets_per_im = pose_targets_per_im.to_object_boxlist()
            assert bbox_targets_per_im.mode == "xyxy"
            bboxes_per_im = bbox_targets_per_im.bbox
            labels_per_im = pose_targets_per_im.class_ids + 1
            anchors_per_im = cat_boxlist(anchors[im_i])
            num_gt = bboxes_per_im.shape[0]
            assert(level_cnt == len(anchors[im_i]))
            #
            rotations_per_im = pose_targets_per_im.rotations
            translations_per_im = pose_targets_per_im.translations
            mask_per_im = pose_targets_per_im.mask

            #
            anchor_sizes_per_level_interest = self.anchor_sizes[:level_cnt]
            anchor_strides_per_level_interst = self.anchor_strides[:level_cnt]
            gt_object_sizes = bbox_targets_per_im.box_span()

            num_anchors_per_level = [len(anchors_per_level.bbox) for anchors_per_level in anchors[im_i]]

            anchors_cx_per_im = (anchors_per_im.bbox[:, 2] + anchors_per_im.bbox[:, 0]) / 2.0
            anchors_cy_per_im = (anchors_per_im.bbox[:, 3] + anchors_per_im.bbox[:, 1]) / 2.0
            anchors_cx_per_im = torch.clamp(anchors_cx_per_im, min = 0, max = mask_per_im.shape[1] - 1).long()
            anchors_cy_per_im = torch.clamp(anchors_cy_per_im, min = 0, max = mask_per_im.shape[0] - 1).long()

            mask_at_anchors = mask_per_im[anchors_cy_per_im, anchors_cx_per_im]
            mask_labels = []
            for gt_i in range(num_gt):
                valid_mask = (mask_at_anchors == (gt_i+1))
                mask_labels.append(valid_mask)
            mask_labels = torch.stack(mask_labels).t()
            mask_labels = mask_labels.long()

            # random selecting candidates from each level first
            candidate_idxs = [[] for i in range(num_gt)]
            start_idx = 0
            gt_sz = gt_object_sizes.view(1,-1).repeat(level_cnt,1)
            lv_sz = torch.FloatTensor(anchor_sizes_per_level_interest).type_as(gt_sz)
            lv_sz = lv_sz.view(-1,1).repeat(1,num_gt)
            dk = torch.log2(gt_sz/lv_sz).abs()
            nk = torch.exp(-self.positive_lambda * (dk * dk))
            nk = self.positive_num * nk / nk.sum(0, keepdim=True)
            nk = (nk + 0.5).int()
            for level in range(level_cnt):
                end_idx = start_idx + num_anchors_per_level[level]
                is_in_mask_per_level = mask_labels[start_idx:end_idx, :]
                #
                for gt_i in range(num_gt):
                    posi_num = nk[level][gt_i]

                    valid_pos = is_in_mask_per_level[:, gt_i].nonzero().view(-1)
                    posi_num = min(posi_num, len(valid_pos))
                    # rand_idx = torch.randint(0, len(valid_pos), (int(posi_num),)) # randoms with replacement
                    rand_idx = torch.randperm(len(valid_pos))[:posi_num] # randoms without replacement
                    candi_pos = valid_pos[rand_idx] + start_idx
                    candidate_idxs[gt_i].append(candi_pos)
                #
                start_idx = end_idx

            # flagging selected positions
            roi = torch.full_like(mask_labels, -INF)
            for gt_i in range(num_gt):
                tmp_idx = torch.cat(candidate_idxs[gt_i], dim=0)
                roi[tmp_idx, gt_i] = 1

            anchors_to_gt_values, anchors_to_gt_indexs = roi.max(dim=1)
            cls_labels_per_im = labels_per_im[anchors_to_gt_indexs]
            cls_labels_per_im[anchors_to_gt_values == -INF] = 0 # background setting

            mask_visibilities, _ = mask_labels.max(dim=1)
            # logical_and, introduced only after pytorch 1.5
            # ignored_indexs = torch.logical_and(mask_visibilities==1, cls_labels_per_im==0)
            ignored_indexs = (mask_visibilities == 1) * (cls_labels_per_im == 0)
            cls_labels_per_im[ignored_indexs] = -1 # positions within mask but not selected will not be touched

            #
            matched_boxes = bboxes_per_im[anchors_to_gt_indexs]
            matched_classes = (labels_per_im - 1)[anchors_to_gt_indexs]
            matched_rotations = rotations_per_im[anchors_to_gt_indexs]
            matched_translations = translations_per_im[anchors_to_gt_indexs]

            #
            matched_3Ds = pose_targets_per_im.keypoints_3d[matched_classes]
            # matched_Ks = pose_targets_per_im.K.repeat(matched_classes.shape[0], 1, 1)
            # TODO
            # assert equals self K
            if not isinstance(self.internal_K, torch.Tensor):
                self.internal_K = torch.FloatTensor(self.internal_K).to(device=matched_3Ds.device).view(3, 3)
            reg_targets_per_im = self.target_coder.encode(
                self.internal_K, matched_3Ds,
                matched_rotations, matched_translations,
                anchors_per_im.bbox)
            cls_labels.append(cls_labels_per_im)
            reg_targets.append(reg_targets_per_im)
            aux_raw_boxes.append(matched_boxes)
            matched_3D_in_camera_frame = torch.bmm(matched_rotations, matched_3Ds.transpose(1, 2)) + matched_translations
            aux_3D_in_camera_frame.append(matched_3D_in_camera_frame.transpose(1, 2))

        return cls_labels, reg_targets, aux_raw_boxes, aux_3D_in_camera_frame

    def __call__(self, pred_cls, pred_reg, targets, anchors):
        labels, reg_targets, aux_raw_boxes, aux_3D_in_camera_frame = self.prepare_targets(targets, anchors)

        N = len(labels)
        pred_cls_flatten, pred_reg_flatten = concat_box_prediction_layers(pred_cls, pred_reg)

        labels_flatten = torch.cat(labels, dim=0)
        reg_targets_flatten = torch.cat(reg_targets, dim=0)
        aux_raw_boxes_flatten = torch.cat(aux_raw_boxes, dim=0)
        aux_3D_in_camera_frame_flatten = torch.cat(aux_3D_in_camera_frame, dim=0)
        anchors_flatten = torch.cat([cat_boxlist(anchors_per_image).bbox for anchors_per_image in anchors], dim=0)

        pos_inds = torch.nonzero(labels_flatten > 0).squeeze(1)

        valid_cls_inds = torch.nonzero(labels_flatten >= 0).squeeze(1)
        cls_loss = self.cls_loss_func(pred_cls_flatten[valid_cls_inds], labels_flatten[valid_cls_inds])

        if pos_inds.numel() > 0:
            pred_reg_flatten = pred_reg_flatten[pos_inds]
            cls_label_flatten = labels_flatten[pos_inds] - 1 # start from class 0
            reg_targets_flatten = reg_targets_flatten[pos_inds]
            aux_raw_boxes_flatten = aux_raw_boxes_flatten[pos_inds]
            aux_3D_in_camera_frame_flatten = aux_3D_in_camera_frame_flatten[pos_inds]
            anchors_flatten = anchors_flatten[pos_inds]

            reg_loss = self.ObjectSpaceLoss(
                pred_reg_flatten, aux_3D_in_camera_frame_flatten,
                cls_label_flatten, anchors_flatten
                )
        else:
            reg_loss = pred_reg_flatten.sum()

        return cls_loss * self.loss_weight_cls, reg_loss * self.loss_weight_reg


# -*- coding: utf-8 -*-
"""
Created on Tue Oct 24 10:46:21 2017

Class to read and manipulate ATL06 data.  Currently set up for Ben-style fake data, should be modified to work with the official ATL06 prodct foramt

@author: ben
"""
import h5py
import numpy as np
from PointDatabase.ATL06_pair import ATL06_pair
from osgeo import osr
import matplotlib.pyplot as plt

class ATL06_data:
    np.seterr(invalid='ignore')
    def __init__(self, beam_pair=2, field_dict=None, list_of_fields=None, NICK=None):
        """
        Initialize an ATl06 data structure.

        ATL06_data has attributes:
            list_of_fields: list of all the fields defined for the current structure
            field_dict: dictionary that shows how to read data from the file.
            beam_pair:  The ICESat-2 beam pair from which the data come
            size: the number of elements in each of its data fields
            shape: the shape of each of its data fields
        It has a variable number of data fields that are assigned during reading,
        each of which contains an  N x 2 numpy array of doubles that holds
        data from the left and right beams of the ATL06 pair (gtxl and gtxr)
        """
        if field_dict is None:
            self.field_dict=self.__default_field_dict__()
        else:
            self.field_dict=field_dict

        if list_of_fields is None:
            list_of_fields=list()
            for group in self.field_dict.keys():
                for field in self.field_dict[group]:
                    list_of_fields.append(field)
        self.filename=None
        self.beam_pair=beam_pair
        self.beam_type=['weak','strong'] # defaults for the early part of the mission
        self.orbit=np.NaN
        self.rgt=np.NaN
        self.file=None
        self.list_of_fields=list_of_fields
        for field in list_of_fields:
            setattr(self, field, np.zeros((0,2)))
        self.size=0
        self.shape=(0,2)

    def __update_size_and_shape__(self):
        """
        When data size and shape may have changed, update the size and shape atttributes
        """
        for field in self.list_of_fields:
            temp=getattr(self, field)
            if hasattr(temp, 'size'):
                self.size=temp.size
                self.shape=temp.shape
                break
        return self

    def __getitem__(self, rows):
        """
        Experimental: index the whole structure at once using square brackets
        """
        return self.copy().subset(rows)

    def __default_field_dict__(self):
        """
        Define the default fields that get read from the h5 file
        """
        field_dict={None:['delta_time','h_li','h_li_sigma','latitude','longitude','atl06_quality_summary','segment_id','sigma_geo_h'],
                    'ground_track':['x_atc', 'y_atc','seg_azimuth','sigma_geo_at','sigma_geo_xt'],
                    'fit_statistics':['dh_fit_dx','dh_fit_dx_sigma','h_mean', 'dh_fit_dy','h_rms_misfit','h_robust_sprd','n_fit_photons', 'signal_selection_source','snr_significance','w_surface_window_final'],
                    'geophysical':['bsnow_conf','bsnow_h','cloud_flg_asr','cloud_flg_atm','r_eff','tide_ocean'],
                    'derived':['valid', 'BP','LR', 'cycle_number', 'rgt']}
        return field_dict

    def from_file(self, filename, index_range=None, x_bounds=None, y_bounds=None):
        """
        Read data from a file.
        """
        self.file=filename
        h5_f=h5py.File(filename,'r')

        beam_names=['gt%d%s' %(self.beam_pair, b) for b in ['l','r']]
        if beam_names[0] not in h5_f or beam_names[1] not in h5_f:
            # return empty data structure
            for group in self.field_dict:
                for field in self.field_dict[group]:
                    setattr(self, field, np.zeros((0,2)))
                self.__update_size_and_shape__()
            return self
        if beam_names[0] not in h5_f.keys():
            return None
        # get the strong/weak beam info
        for count, beam_name in enumerate(beam_names):
            try:
                self.beam_type[count]=h5_f[beam_name]['atlas_beam_type']
            except KeyError:
                pass  # leave the beam type as the default
        if index_range is None or index_range[1]==-1:
            index_range=[0, np.minimum(h5_f[beam_names[0]]['land_ice_segments']['h_li'].size, h5_f[beam_names[1]]['land_ice_segments']['h_li'].size)]
        n_vals=index_range[-1]-index_range[0]
        # read the orbit number
        try:
            self.rgt=int(h5_f['orbit_info']['rgt'][0])+np.zeros((n_vals,2))
            self.orbit=int(h5_f['orbit_info']['orbit_number'][0])
        except:
            pass

        for group in self.field_dict.keys():
            for field in self.field_dict[group]:
                if field not in self.list_of_fields:
                    self.list_of_fields.append(field)
                try:
                    #bad_val=-9999 # default value, works for most groups
                    if group is None:
                        #The None group corresponds to the top level of the land_ice_segments hirearchy
                        # most datasets have a __fillValue attribute marking good data, but some don't
                        try:
                            bad_val=h5_f[beam_names[0]]['land_ice_segments'][field].attrs['_FillValue']
                        except KeyError:
                            #latitude and longitude don't have fill values, but -9999 works OK.
                            bad_val=-9999
                        data=np.c_[
                            np.array(h5_f[beam_names[0]]['land_ice_segments'][field][index_range[0]:index_range[1]]).transpose(),
                            np.array(h5_f[beam_names[1]]['land_ice_segments'][field][index_range[0]:index_range[1]]).transpose()]
                    elif group == "orbit_info":
                        data=np.zeros((n_vals, 2))+h5_f['orbit_info'][field]
                        bad_val=-9999
                    elif group == "derived" and field == "valid":
                        data=np.ones((index_range[1]-index_range[0], 2), dtype='bool')
                        bad_val=0
                    elif group == "derived" and field == "BP":
                        data=np.zeros((index_range[1]-index_range[0], 2))+self.beam_pair
                        bad_val=0
                    elif group == "derived" and field == "LR":
                        data=np.ones((index_range[1]-index_range[0], 2))
                        data[:,0]=0
                        bad_val=-9999
                    elif group == "derived" and field == "n_pixels":
                        if self.beam_type[0]=='weak':
                            data=np.tile([4, 16], [n_vals, 1])
                        else:
                            data=np.tile([16, 4], [n_vals, 1])
                        bad_val=-9999
                    elif group == "derived" and field == "spot":
                        data=np.ones((index_range[1]-index_range[0], 2))
                        for bb in [0, 1]:
                            data[:,bb]=np.float64(h5_f[beam_names[bb]].attrs['atlas_spot_number'])
                    elif group == "derived" and field == "sigma_geo_r":
                        # Earlier versions of the data don't include a sigma_geo_r.  This is a fake value
                        data=np.zeros((index_range[1]-index_range[0], 2)) +0.03
                    elif group == "derived":
                        continue
                    else:
                        # All other groups are under the land_ice_segments/group hirearchy
                         try:
                            bad_val=h5_f[beam_names[0]]['land_ice_segments'][group][field].attrs['_FillValue']
                         except KeyError:
                            #flags don't have fill values, but -9999 works OK.
                            bad_val=-9999
                         try:
                             data=np.c_[
                                np.array(h5_f[beam_names[0]]['land_ice_segments'][group][field][index_range[0]:index_range[1]]).transpose(),
                                np.array(h5_f[beam_names[1]]['land_ice_segments'][group][field][index_range[0]:index_range[1]]).transpose()]
                         except KeyError:
                             print("missing hdf field for land_ice_segments/%s/%s in %s" % (group, field, filename))
                             data=np.zeros((index_range[1]+1-index_range[0], 2))+bad_val
                    # identify the bad data elements before converting the field to double
                    bad=data==bad_val
                    if field != 'valid':
                        data=data.astype(np.float64)
                    # mark data that are invalid (according to the h5 file) with NaNs
                    data[bad]=np.NaN
                    setattr(self, field, data)
                except KeyError:
                    print("could not read %s/%s" % (group, field))
                    setattr(self, field, np.zeros( [n_vals, 2])+np.NaN)
#        if 'atl06_quality_summary' in self.list_of_fields:
#            # add a non-broken version of the atl06 quality summary
#            setattr(self, 'atl06_quality_summary', (self.h_li_sigma > 1) | (self.h_robust_sprd > 1) | (self.snr_significance > 0.02))
        if "derived" in self.field_dict and "matlab_time" in self.field_dict['derived']:
            self.get_Matlab_time()
        self.__update_size_and_shape__()

        if "derived" in self.field_dict and "rss_along_track_dh" in self.field_dict['derived']:
            self.get_rss_along_track_dh()

        if "derived" in self.field_dict and "min_along_track_dh" in self.field_dict['derived']:
            self.get_min_along_track_dh()
        # assign fields that must be copied from single-value attributes in the
        # h5 file
        if 'cycle_number' in self.list_of_fields:
            # get the cycle number
            try:
                cycle_number=h5_f['ancillary_data']['start_cycle']
            except KeyError:
                cycle_number=-9999
            setattr(self, 'cycle_number', cycle_number+np.zeros(self.shape, dtype=np.float64))
        h5_f.close()
        return self

    def get_xy(self, proj4_string=None, EPSG=None):
        # method to get projected coordinates for the data.  Adds 'x' and 'y' fields to the data, optionally returns 'self'
        out_srs=osr.SpatialReference()
        if proj4_string is None and EPSG is not None:
            out_srs.ImportFromEPSG(EPSG)
        else:
            projError= out_srs.ImportFromProj4(proj4_string)
            if projError > 0:
                out_srs.ImportFromWkt(proj4_string)
        ll_srs=osr.SpatialReference()
        ll_srs.ImportFromEPSG(4326)
        if hasattr(osr,'OAMS_TRADITIONAL_GIS_ORDER'):
            ll_srs.SetAxisMappingStrategy(osr.OAMS_TRADITIONAL_GIS_ORDER)
        ct=osr.CoordinateTransformation(ll_srs, out_srs)
        if self.latitude.size==0:
            self.x=np.zeros_like(self.latitude)
            self.y=np.zeros_like(self.latitude)
        else:
            x, y, z= list(zip(*[ct.TransformPoint(*xyz) for xyz in zip(np.ravel(self.longitude), np.ravel(self.latitude), np.zeros_like(np.ravel(self.latitude)))]))
            self.x=np.reshape(x, self.latitude.shape)
            self.y=np.reshape(y, self.longitude.shape)
        if 'x' not in self.list_of_fields:
            self.list_of_fields += ['x','y']
        return self

    def append(self, D):
        """
        Append the fields of two ATL06_data instances
        """
        for field in self.list_of_fields:
            setattr(self, np.c_[getattr(self, field), getattr(D, field)])
        self.__update_size_and_shape__()
        return self

    def assign(self, D):
        """
        Assign fields from a dictionary to self
        """
        for key in D:
            setattr(self, key, D[key])
            if key not in self.list_of_fields:
                self.list_of_fields.append(key)
        return self

    def ravel_fields(self):
        """
        ravel all the fields in self
        """
        for field in self.list_of_fields:
            setattr(self, field, getattr(self, field).ravel())
        self.__update_size_and_shape__()
        return self

    def from_list(self, D6_list, copy_fields=True):
        """
        Append the fields of several ATL06_data instances.
        """
        if copy_fields:
            fields=[]
            for D6 in D6_list:
                try:
                    fields += D6.list_of_fields
                except:
                    pass
            self.list_of_fields=list(set(fields))
        try:
            for field in self.list_of_fields:
                data_list=[getattr(this_D6, field) for this_D6 in D6_list]
                setattr(self, field, np.concatenate(data_list, 0))
        except TypeError:
            for field in self.list_of_fields:
                setattr(self, field, getattr(D6_list, field))
        except ValueError:
            print("ValueError")
        self.__update_size_and_shape__()
        return self

    def from_dict(self, D6_dict):
        """
        Build an ATL06_data from a dictionary
        """
        for field in self.list_of_fields:
            try:
                setattr(self, field, D6_dict[field])
            except KeyError:
                setattr(self, field, None)
        self.__update_size_and_shape__()
        return self

    def index(self, rows):
        """
        Select a subset of rows within a given ATL06_data instance (modify in place)
        """
        for field in self.list_of_fields:
            if len(self.shape) > 1:
                setattr(self, field, getattr(self, field)[rows,:])
            else:
                setattr(self, field, getattr(self, field)[rows])
        self.__update_size_and_shape__()
        return self

    def subset(self, index, by_row=True, datasets=None):
        """
        Make a subsetted copy of the current ATL06_data instance
        """
        dd=dict()
        if datasets is None:
            datasets=self.list_of_fields
        for field in datasets:
            try:
                if by_row is not None and by_row:
                    dd[field]=getattr(self, field)[index,:]
                else:
                    dd[field]=getattr(self, field)[index,:].ravel()[index]
            except IndexError:
                print("IndexError")
        return ATL06_data(list_of_fields=datasets, beam_pair=self.beam_pair).from_dict(dd)

    def copy(self):
        return ATL06_data( list_of_fields=self.list_of_fields, beam_pair=self.beam_pair).from_list((self))

    def plot(self, valid_pairs=None, valid_segs=None):
        colors=('r','b')
        for col in (0, 1):
            #plt.errorbar(self.x_atc[:,col], self.h_li[:,col], yerr=self.h_li_sigma[:, col], c=colors[col], marker='.', linestyle='None', markersize=4);
            plt.plot(self.x_atc[:,col], self.h_li[:,col], c=colors[col], marker='.', linestyle='None', markersize=2);
        if valid_segs is not None:
            for col in (0, 1):
                plt.plot(self.x_atc[valid_segs[col], col], self.h_li[valid_segs[col], col], marker='x',c=colors[col])

        if valid_pairs is not None:
            for col in (0, 1):
                plt.plot(self.x_atc[valid_pairs, col], self.h_li[valid_pairs, col], marker='o',c=colors[col])
        plt.ylim(np.nanmin(self.h_li[self.atl06_quality_summary[:,1]==0,1])-5., np.nanmax(self.h_li[self.atl06_quality_summary[:,1]==0 ,1])+5 )
        #plt.show()
        return

    def get_pairs(self, datasets=None):
        pair_list=list()
        for i in np.arange(self.h_li.shape[0]):
            pair_list.append(ATL06_pair(D6=self.subset(i, by_row=True, datasets=datasets)))
        all_pairs=ATL06_pair(pair_data=pair_list)
        return all_pairs

    def get_Matlab_time(self):
        self.assign({'matlab_time': 737061. + self.delta_time/24./3600.})

    def get_rss_along_track_dh(self):
        self.rss_along_track_dh=np.zeros(self.shape)
        n_pts=self.shape[0]
        if n_pts > 1:
            i0=slice(1, n_pts-1)
            for ii in [-1, 1]:
                i1=slice(1+ii, n_pts-1+ii)
                dx=self.x_atc[i0,:]-self.x_atc[i1,:]
                self.rss_along_track_dh[i0,:] += (self.h_li[i0,:]-self.dh_fit_dx[i0,:]*dx-self.h_li[i1,:])**2
            self.rss_along_track_dh[0,:]=(self.h_li[1,:]-self.h_li[0,:] - (self.x_atc[1,:]-self.x_atc[0,:])*self.dh_fit_dx[0,:])**2
            self.rss_along_track_dh[-1,:]=(self.h_li[-1,:]-self.h_li[-2,:] - (self.x_atc[-1,:]-self.x_atc[-2,:])*self.dh_fit_dx[-1,:])**2
            self.rss_along_track_dh = np.sqrt(self.rss_along_track_dh)

    def get_min_along_track_dh(self):
        self.min_along_track_dh=np.zeros(list(self.shape)+[2])
        n_pts=self.shape[0]
        if n_pts > 1:
            i0=slice(1, n_pts-1)
            for dim3, ii in enumerate([-1, 1]):
                i1=slice(1+ii, n_pts-1+ii)
                dx=self.x_atc[i0,:]-self.x_atc[i1,:]
                self.min_along_track_dh[i0,:, dim3] = np.abs(self.h_li[i0,:]-self.dh_fit_dx[i0,:]*dx-self.h_li[i1,:])
            self.min_along_track_dh = np.nanmin(self.min_along_track_dh, axis=2)
            self.min_along_track_dh[0,:]=np.abs(self.h_li[1,:]-self.h_li[0,:] - (self.x_atc[1,:]-self.x_atc[0,:])*self.dh_fit_dx[0,:])
            self.min_along_track_dh[-1,:]=np.abs(self.h_li[-1,:]-self.h_li[-2,:] - (self.x_atc[-1,:]-self.x_atc[-2,:])*self.dh_fit_dx[-1,:])
        else:
            self.min_along_track_dh=np.zeros([1,2])+np.NaN

def delta_t_to_Matlab(delta_t):
    return 730486 + delta_t/24./3600.


# -*- coding: utf-8 -*-
"""
Created on Sat Mar 14 19:13:15 2020

@author: Manuel Camargo
"""
import itertools
import numpy as np
# import pandas as pd

from nltk.util import ngrams
import keras.utils as ku

# from support_modules import nn_support as nsup
# import support as sup
# import pandas as pd
# import json

class SequencesCreator():

    def __init__(self, log, ac_index, rl_index):
        """constructor"""
        self.log = log
        self.ac_index = ac_index
        self.rl_index = rl_index

    def vectorize(self, model_type, params):
        loader = self._get_vectorizer(model_type)
        return loader(params)

    def _get_vectorizer(self, model_type):
        if model_type == 'shared_cat':
            return self._vectorize_seq
        elif model_type == 'shared_cat_inter':
            return self._vectorize_seq_inter
        elif model_type == 'shared_cat_rd':
            return self._vectorize_seq_rd
        elif model_type == 'shared_cat_wl':
            return self._vectorize_seq_wl
        elif model_type == 'shared_cat_cx':
            return self._vectorize_seq_cx
        elif model_type == 'cnn_lstm':
            return self._vectorize_seq
        elif model_type == 'cnn_lstm_inter':
            return self._vectorize_seq_inter
        elif model_type == 'seq2seq':
            return self._vectorize_seq2seq
        elif model_type == 'seq2seq_inter':
            return self._vectorize_seq2seq_inter
        elif model_type == 'shared_cat_inter_full':
            return self._vectorize_seq_inter_full
        elif model_type == 'cnn_lstm_inter_full':
            return self._vectorize_seq_inter_full
        elif model_type == 'shared_cat_city':
            return self._vectorize_seq_city
        elif model_type == 'shared_cat_snap':
            return self._vectorize_seq_snap
        else:
            raise ValueError(model_type)

    def _vectorize_seq(self, parms):
        """
        Example function with types documented in the docstring.
        parms:
            log_df (dataframe): event log data.
            ac_index (dict): index of activities.
            rl_index (dict): index of roles.
            parms (dict): parms for training the network
        Returns:
            dict: Dictionary that contains all the LSTM inputs.
        """
        equi = {'ac_index': 'activities',
                'rl_index': 'roles',
                'dur_norm': 'times'}
        columns = list(equi.keys())
        vec = {'prefixes': dict(),
               'next_evt': dict()}
        self.log = self.reformat_events(columns, parms['one_timestamp'])
        # n-gram definition
        for i, _ in enumerate(self.log):
            for x in columns:
                serie = list(ngrams(self.log[i][x], parms['n_size'],
                                    pad_left=True, left_pad_symbol=0))
                y_serie = [x[-1] for x in serie]
                serie = serie[:-1]
                y_serie = y_serie[1:]
                vec['prefixes'][equi[x]] = (vec['prefixes'][equi[x]] + serie
                                            if i > 0 else serie)
                vec['next_evt'][equi[x]] = (vec['next_evt'][equi[x]] + y_serie
                                            if i > 0 else y_serie)

        # Transform task, dur and role prefixes in vectors
        for value in equi.values():
            vec['prefixes'][value] = np.array(vec['prefixes'][value])
            vec['next_evt'][value] = np.array(vec['next_evt'][value])
        # Reshape dur (prefixes, n-gram size, 1) i.e. time distribute
        vec['prefixes']['times'] = vec['prefixes']['times'].reshape(
                (vec['prefixes']['times'].shape[0],
                 vec['prefixes']['times'].shape[1], 1))
        # one-hot encode target values
        vec['next_evt']['activities'] = ku.to_categorical(
            vec['next_evt']['activities'], num_classes=len(self.ac_index))
        vec['next_evt']['roles'] = ku.to_categorical(
            vec['next_evt']['roles'], num_classes=len(self.rl_index))
        return vec

    def _vectorize_seq_inter(self, parms):
        # columns to keep
        columns = ['ev_rd_norm', 'ev_rp_occ_norm', 'ev_et_norm', 'ev_et_t_norm',
                   'ac_index', 'rl_index', 'dur_norm']
        return self.process_intercases(columns, parms)

    def _vectorize_seq_inter_full(self, parms):
        # columns to keep
        columns = ['acc_cycle_norm', 'daytime_norm', 'ev_rd_norm',
                   'ev_rp_occ_norm', 'ev_et_norm', 'ev_et_t_norm',
                   'ac_index', 'rl_index', 'dur_norm']
        return self.process_intercases(columns, parms)

    def _vectorize_seq_rd(self, parms):
        # columns to keep
        columns = ['ev_rd_norm', 'ev_rp_occ_norm',
                   'ac_index', 'rl_index', 'dur_norm']
        return self.process_intercases(columns, parms)

    def _vectorize_seq_wl(self, parms):
        # columns to keep
        columns = ['ev_et_norm', 'ev_et_t_norm',
                   'ac_index', 'rl_index', 'dur_norm']
        return self.process_intercases(columns, parms)

    def _vectorize_seq_cx(self, parms):
        # columns to keep
        columns = ['acc_cycle_norm', 'daytime_norm',
                   'ac_index', 'rl_index', 'dur_norm']
        return self.process_intercases(columns, parms)

    def _vectorize_seq_city(self, parms):
        # columns to keep
        columns = ['city1_norm','city2_norm','city3_norm',
                   'ac_index', 'rl_index', 'dur_norm']
        return self.process_intercases(columns, parms)

    def _vectorize_seq_snap(self, parms):
        # columns to keep
        columns = ['snap1_norm','snap2_norm','snap3_norm',
                   'ac_index', 'rl_index', 'dur_norm']
        return self.process_intercases(columns, parms)

    def process_intercases(self, columns, parms):
        vec = {'prefixes': dict(),
               'next_evt': dict()}
        self.log = self.reformat_events(columns, parms['one_timestamp'])
        # n-gram definition
        equi = {'ac_index': 'activities',
                'rl_index': 'roles',
                'dur_norm': 'times'}
        x_inter_dict = dict()
        y_inter_dict = dict()
        for i, _ in enumerate(self.log):
            for x in columns:
                serie = list(ngrams(self.log[i][x], parms['n_size'],
                                    pad_left=True, left_pad_symbol=0))
                y_serie = [x[-1] for x in serie]
                serie = serie[:-1]
                y_serie = y_serie[1:]
                if x in list(equi.keys()):
                    vec['prefixes'][equi[x]] = (
                        vec['prefixes'][equi[x]] + serie if i > 0 else serie)
                    vec['next_evt'][equi[x]] = (
                        vec['next_evt'][equi[x]] + y_serie
                        if i > 0 else y_serie)
                else:
                    x_inter_dict[x] = (
                        x_inter_dict[x] + serie if i > 0 else serie)
                    y_inter_dict[x] = (
                        y_inter_dict[x] + y_serie if i > 0 else y_serie)
        # Transform task, dur and role prefixes in vectors
        for value in equi.values():
            vec['prefixes'][value] = np.array(vec['prefixes'][value])
            vec['next_evt'][value] = np.array(vec['next_evt'][value])
        # Reshape dur (prefixes, n-gram size, 1) i.e. time distribute
        vec['prefixes']['times'] = vec['prefixes']['times'].reshape(
                (vec['prefixes']['times'].shape[0],
                  vec['prefixes']['times'].shape[1], 1))
        # one-hot encode target values
        vec['next_evt']['activities'] = ku.to_categorical(
            vec['next_evt']['activities'], num_classes=len(self.ac_index))
        vec['next_evt']['roles'] = ku.to_categorical(
            vec['next_evt']['roles'], num_classes=len(self.rl_index))
        # Reshape intercase attributes (prefixes, n-gram size, number of attributes)
        for key, value in x_inter_dict.items():
            x_inter_dict[key] = np.array(value)
            x_inter_dict[key] = x_inter_dict[key].reshape(
                (x_inter_dict[key].shape[0], x_inter_dict[key].shape[1], 1))
        vec['prefixes']['inter_attr'] = np.dstack(list(x_inter_dict.values()))
        # Reshape y intercase attributes (suffixes, number of attributes)
        for key, value in y_inter_dict.items():
            x_inter_dict[key] = np.array(value)
        vec['next_evt']['inter_attr'] = np.dstack(list(y_inter_dict.values()))[0]
        return vec

    def _vectorize_seq2seq(self, parms):
        """Example function with types documented in the docstring.
        parms:
            log_df (dataframe): event log data.
            ac_index (dict): index of activities.
            rl_index (dict): index of roles.
            parms (dict): parms for training the network
        Returns:
            dict: Dictionary that contains all the LSTM inputs.
        """
        columns = ['ac_index', 'rl_index', 'dur_norm']
        examples = {'encoder_input_data': dict(),
                    'decoder_input_data': dict(),
                    'decoder_target_data': dict()}
        self.log = self.reformat_events(columns, parms['one_timestamp'])
        max_length = np.max([len(x['ac_index']) for x in self.log])
        # n-gram definition
        equi = {'ac_index': 'activities',
                'rl_index': 'roles',
                'dur_norm': 'times'}
        for i, _ in enumerate(self.log):
            for x in columns:
                serie_e, serie_d, serie_dt = list(), list(), list()
                for idx in range(1, len(self.log[i][x])):
                    serie_e.append(
                        [0]*(max_length - idx) + self.log[i][x][:idx])
                    serie_d.append(
                        self.log[i][x][idx:] +
                        [0]*(max_length - len(self.log[i][x][idx:])))
                    serie_dt.append(
                        self.log[i][x][idx+1:] +
                        [0]*((max_length - len(self.log[i][x][idx:]))+1))
                examples['encoder_input_data'][equi[x]] = (
                    examples['encoder_input_data'][equi[x]] + serie_e
                    if i > 0 else serie_e)
                examples['decoder_input_data'][equi[x]] = (
                    examples['decoder_input_data'][equi[x]] + serie_d
                    if i > 0 else serie_d)
                examples['decoder_target_data'][equi[x]] = (
                    examples['decoder_target_data'][equi[x]] + serie_dt
                    if i > 0 else serie_dt)
        for value in equi.values():
            examples['encoder_input_data'][value]= np.array(
                examples['encoder_input_data'][value])
            examples['decoder_input_data'][value]= np.array(
                examples['decoder_input_data'][value])
            examples['decoder_target_data'][value]= np.array(
                examples['decoder_target_data'][value])
        # Reshape dur (prefixes, n-gram size, 1) i.e. time distribute
        examples['encoder_input_data']['times'] = (
            examples['encoder_input_data']['times'].reshape(
                (examples['encoder_input_data']['times'].shape[0],
                 examples['encoder_input_data']['times'].shape[1], 1)))
        examples['decoder_input_data']['times'] = (
            examples['decoder_input_data']['times'].reshape(
                (examples['decoder_input_data']['times'].shape[0],
                 examples['decoder_input_data']['times'].shape[1], 1)))
        examples['decoder_target_data']['times'] = (
            examples['decoder_target_data']['times'].reshape(
                (examples['decoder_target_data']['times'].shape[0],
                 examples['decoder_target_data']['times'].shape[1], 1)))
        # One hot encode decoder_target_data
        examples['decoder_target_data']['activities'] = ku.to_categorical(
            examples['decoder_target_data']['activities'],
            num_classes=len(self.ac_index))
        examples['decoder_target_data']['roles'] = ku.to_categorical(
            examples['decoder_target_data']['roles'],
            num_classes=len(self.rl_index))
        return examples

    def _vectorize_seq2seq_inter(self, parms):
        """
        Example function with types documented in the docstring.
        parms:
            log_df (dataframe): event log data.
            ac_index (dict): index of activities.
            rl_index (dict): index of roles.
            parms (dict): parms for training the network
        Returns:
            dict: Dictionary that contains all the LSTM inputs.
        """
    #    log_df = log_df[log_df.caseid=='Case28']
        columns = ['caseid', 'task', 'user', 'start_timestamp', 'end_timestamp',
                   'dur_log', 'role', 'event_id', 'ev_duration', 'dur', 'ev_rd']
        columns = [x for x in list(self.log.columns) if x not in columns]
        examples = {'encoder_input_data': dict(),
                    'decoder_target_data': dict()}
        self.log = self.reformat_events(columns, parms['one_timestamp'])
        max_length = np.max([len(x['ac_index']) for x in
                             self.reformat_events(['ac_index'],
                                                  parms['one_timestamp'])])
        # n-gram definition
        equi = {'ac_index':'activities', 'rl_index':'roles', 'dur_norm':'times'}
        x_inter_dict = dict()
        for i, _ in enumerate(self.log):
            for x in columns:
                serie, y_serie = list(), list()
                for idx in range(1, len(self.log[i][x])):
                    serie.append([0]*(max_length - idx) + self.log[i][x][:idx])
                    y_serie.append(self.log[i][x][idx:] + [0]*(max_length - len(self.log[i][x][idx:])))
                if x in list(equi.keys()):
                    examples['encoder_input_data'][equi[x]] = examples['encoder_input_data'][equi[x]] + serie if i > 0 else serie
                    examples['decoder_target_data'][equi[x]] = examples['decoder_target_data'][equi[x]] + y_serie if i > 0 else y_serie
                else:
                    x_inter_dict[x] = x_inter_dict[x] + serie if i > 0 else serie
        # Transform task, dur and role prefixes in vectors
        for value in equi.values():
            examples['encoder_input_data'][value]= np.array(examples['encoder_input_data'][value])
            examples['decoder_target_data'][value]= np.array(examples['decoder_target_data'][value])
        # Reshape dur (prefixes, n-gram size, 1) i.e. time distribute
        examples['encoder_input_data']['times'] = examples['encoder_input_data']['times'].reshape(
                (examples['encoder_input_data']['times'].shape[0],
                 examples['encoder_input_data']['times'].shape[1], 1))
        examples['decoder_target_data']['times'] = examples['decoder_target_data']['times'].reshape(
                (examples['decoder_target_data']['times'].shape[0],
                 examples['decoder_target_data']['times'].shape[1], 1))

        # One hot encode decoder_target_data
        examples['decoder_target_data']['activities'] = ku.to_categorical(examples['decoder_target_data']['activities'],
                                                        num_classes=len(self.ac_index))
        examples['decoder_target_data']['roles'] = ku.to_categorical(examples['decoder_target_data']['roles'],
                                                        num_classes=len(self.rl_index))
        # Reshape intercase attributes (prefixes, n-gram size, number of attributes)
        for key, value in x_inter_dict.items():
            x_inter_dict[key] = np.array(value)
            x_inter_dict[key] = x_inter_dict[key].reshape((x_inter_dict[key].shape[0],
                       x_inter_dict[key].shape[1], 1))
        examples['encoder_input_data']['inter_attr'] = np.dstack(list(x_inter_dict.values()))
    #    examples['decoder_target_data']['inter_attr'] = np.dstack(list(y_inter_dict.values()))[0]
        return examples

    # =============================================================================
    # Reformat events
    # =============================================================================
    def reformat_events(self, columns, one_timestamp):
        """Creates series of activities, roles and relative times per trace.
        parms:
            self.log: dataframe.
            ac_index (dict): index of activities.
            rl_index (dict): index of roles.
        Returns:
            list: lists of activities, roles and relative times.
        """
        temp_data = list()
        log_df = self.log.to_dict('records')
        key = 'end_timestamp' if one_timestamp else 'start_timestamp'
        log_df = sorted(log_df, key=lambda x: (x['caseid'], key))
        for key, group in itertools.groupby(log_df, key=lambda x: x['caseid']):
            trace = list(group)
            temp_dict = dict()
            for x in columns:
                serie = [y[x] for y in trace]
                if x == 'ac_index':
                    serie.insert(0, self.ac_index[('start')])
                    serie.append(self.ac_index[('end')])
                elif x == 'rl_index':
                    serie.insert(0, self.rl_index[('start')])
                    serie.append(self.rl_index[('end')])
                else:
                    serie.insert(0, 0)
                    serie.append(0)
                temp_dict = {**{x: serie}, **temp_dict}
            temp_dict = {**{'caseid': key}, **temp_dict}
            temp_data.append(temp_dict)
        return temp_data


"""
StyleMapGAN
Copyright (c) 2021-present NAVER Corp.

This work is licensed under the Creative Commons Attribution-NonCommercial
4.0 International License. To view a copy of this license, visit
http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to
Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
"""

import argparse
import pickle
import torch
from torch import nn
from torch.nn import functional as F
from torch.utils import data
from torchvision import utils, transforms
import numpy as np
from torchvision.datasets import ImageFolder
from training.dataset import (
    MultiResolutionDataset,
    GTMaskDataset,
)
from scipy import linalg
import random
import time
import os
from tqdm import tqdm
from copy import deepcopy
import cv2
from PIL import Image
from itertools import combinations
from training.model import Generator, Encoder

random.seed(0)
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)

def save_image(img, path, normalize=True, range=(-1, 1)):
    utils.save_image(
        img,
        path,
        normalize=normalize,
        range=range,
    )

def save_images(imgs, paths, normalize=True, range=(-1, 1)):
    for img, path in zip(imgs, paths):
        save_image(img, path, normalize=normalize, range=range)

def make_noise(batch, latent_channel_size, device):
    return torch.randn(batch, latent_channel_size, device=device)

def data_sampler(dataset, shuffle):
    if shuffle:
        return data.RandomSampler(dataset)
    else:
        return data.SequentialSampler(dataset)

class Model(nn.Module):
    def __init__(self, device="cuda"):
        super(Model, self).__init__()
        self.g_ema = Generator(
            args.size,
            args.mapping_layer_num,
            args.latent_channel_size,
            args.latent_spatial_size,
            lr_mul=args.lr_mul,
            channel_multiplier=args.channel_multiplier,
            normalize_mode=args.normalize_mode,
            small_generator=args.small_generator,
        )
        self.e_ema = Encoder(
            args.size,
            args.latent_channel_size,
            args.latent_spatial_size,
            channel_multiplier=args.channel_multiplier,
        )

    def forward(self, input, mode):
        if mode == "projection":
            fake_stylecode = self.e_ema(input)

            return fake_stylecode

        elif mode == "calculate_mean_stylemap":
            truncation_mean_latent = self.g_ema(input, calculate_mean_stylemap=True)

            return truncation_mean_latent

        elif mode == "random_generation":
            z, truncation, truncation_mean_latent = input

            fake_img, _ = self.g_ema(
                z,
                truncation=truncation,
                truncation_mean_latent=truncation_mean_latent,
            )

            return fake_img

        elif mode == "reconstruction":
            fake_stylecode = self.e_ema(input)
            fake_img, _ = self.g_ema(fake_stylecode, input_is_stylecode=True)
            return fake_img

        elif mode == "stylemixing":
            w1, w2 = input
            fake_imgs = torch.Tensor().to(device)

            coarse_img, _ = self.g_ema(
                [w1, w2],
                input_is_stylecode=True,
                mix_space=f"stylemixing_coarse",
            )
            fine_img, _ = self.g_ema(
                [w1, w2], input_is_stylecode=True, mix_space=f"stylemixing_fine"
            )

            return coarse_img, fine_img

        elif mode == "w_interpolation":
            w1, w2 = input
            lambda_w = random.random()
            w = w1 * lambda_w + w2 * (1 - lambda_w)
            w = w.unsqueeze(0)
            fake_img, _ = self.g_ema(w, input_is_stylecode=True)

            return fake_img

        elif mode == "local_editing":
            w1, w2, mask = input
            w1, w2, mask = w1.unsqueeze(0), w2.unsqueeze(0), mask.unsqueeze(0)

            if dataset_name == "celeba_hq":
                mixed_image = self.g_ema(
                    [w1, w2],
                    input_is_stylecode=True,
                    mix_space="w_plus",
                    mask=mask,
                )[0]

            elif dataset_name == "afhq":
                mixed_image = self.g_ema(
                    [w1, w2], input_is_stylecode=True, mix_space="w", mask=mask
                )[0]

            recon_img_src, _ = self.g_ema(w1, input_is_stylecode=True)
            recon_img_ref, _ = self.g_ema(w2, input_is_stylecode=True)

            return mixed_image, recon_img_src, recon_img_ref

        elif mode == "transplantation":
            src_img, ref_img, coordinates = input

            src_img, ref_img = (
                src_img.unsqueeze(0),
                ref_img.unsqueeze(0),
            )
            src_w = self.e_ema(src_img)
            ref_w = self.e_ema(ref_img)
            recon_img_src, _ = self.g_ema(src_w, input_is_stylecode=True)
            recon_img_ref, _ = self.g_ema(ref_w, input_is_stylecode=True)

            for (
                (src_p_y, src_p_x),
                (ref_p_y, ref_p_x),
                height,
                width,
            ) in coordinates:
                mask_src = -torch.ones([8, 8]).to(device)
                mask_ref = -torch.ones([8, 8]).to(device)

                mask_src[src_p_y : src_p_y + height, src_p_x : src_p_x + width] = 1
                mask_ref[ref_p_y : ref_p_y + height, ref_p_x : ref_p_x + width] = 1

                mask_src, mask_ref = mask_src.unsqueeze(0), mask_ref.unsqueeze(0)
                mask_src = mask_src.unsqueeze(1).repeat(1, 64, 1, 1)
                mask_ref = mask_ref.unsqueeze(1).repeat(1, 64, 1, 1)
                src_w[mask_src == 1] = ref_w[mask_ref == 1]

            mixed_image, _ = self.g_ema(src_w, input_is_stylecode=True)

            return mixed_image, recon_img_src, recon_img_ref

if __name__ == "__main__":
    device = "cuda"

    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--mixing_type",
        choices=[
            "local_editing",
            "transplantation",
            "w_interpolation",
            "reconstruction",
            "stylemixing",
            "random_generation",
        ],
        required=True,
    )
    parser.add_argument("--ckpt", metavar="CHECKPOINT", required=True)
    parser.add_argument("--test_lmdb", type=str)
    parser.add_argument("--batch", type=int, default=16)
    parser.add_argument("--num_workers", type=int, default=10)
    parser.add_argument("--save_image_dir", type=str, default="expr")

    # Below argument is needed for local editing.
    parser.add_argument(
        "--local_editing_part",
        type=str,
        default=None,
        choices=[
            "nose",
            "hair",
            "background",
            "eye",
            "eyebrow",
            "lip",
            "neck",
            "cloth",
            "skin",
            "ear",
        ],
    )

    args = parser.parse_args()

    ckpt = torch.load(args.ckpt)
    train_args = ckpt["train_args"]
    for key in vars(train_args):
        if not (key in vars(args)):
            setattr(args, key, getattr(train_args, key))
    print(args)

    dataset_name = args.dataset
    args.save_image_dir = os.path.join(
        args.save_image_dir, args.mixing_type, dataset_name
    )

    model = Model().to(device)
    model.g_ema.load_state_dict(ckpt["g_ema"])
    model.e_ema.load_state_dict(ckpt["e_ema"])
    model.eval()

    batch = args.batch

    device = "cuda"
    transform = transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True),
        ]
    )

    if args.mixing_type == "random_generation":
        os.makedirs(args.save_image_dir, exist_ok=True)
    elif args.mixing_type in [
        "w_interpolation",
        "reconstruction",
        "transplantation",
        "stylemixing",
    ]:
        os.makedirs(args.save_image_dir, exist_ok=True)
        dataset = MultiResolutionDataset(args.test_lmdb, transform, args.size)
    elif args.mixing_type == "local_editing":

        if dataset_name == "afhq":
            args.save_image_dir = os.path.join(args.save_image_dir)
            for kind in [
                "mask",
                "source_image",
                "source_reconstruction",
                "reference_image",
                "reference_reconstruction",
                "synthesized_image",
            ]:
                os.makedirs(os.path.join(args.save_image_dir, kind), exist_ok=True)
        else:  # celeba_hq
            args.save_image_dir = os.path.join(
                args.save_image_dir,
                args.local_editing_part,
            )
            for kind in [
                "mask",
                "mask_ref",
                "mask_src",
                "source_image",
                "source_reconstruction",
                "reference_image",
                "reference_reconstruction",
                "synthesized_image",
            ]:
                os.makedirs(os.path.join(args.save_image_dir, kind), exist_ok=True)
            mask_path_base = f"data/{dataset_name}/local_editing"

        # GT celeba_hq mask images
        if dataset_name == "celeba_hq":
            assert "celeba_hq" in args.test_lmdb

            dataset = GTMaskDataset("data/celeba_hq", transform, args.size)

            parts_index = {
                "background": [0],
                "skin": [1],
                "eyebrow": [6, 7],
                "eye": [3, 4, 5],
                "ear": [8, 9, 15],
                "nose": [2],
                "lip": [10, 11, 12],
                "neck": [16, 17],
                "cloth": [18],
                "hair": [13, 14],
            }

        # afhq, coarse(half-and-half) masks
        else:
            assert "afhq" in args.test_lmdb and "afhq" == dataset_name
            dataset = MultiResolutionDataset(args.test_lmdb, transform, args.size)

    if args.mixing_type in [
        "w_interpolation",
        "reconstruction",
        "stylemixing",
        "local_editing",
    ]:
        n_sample = len(dataset)
        sampler = data_sampler(dataset, shuffle=False)

        loader = data.DataLoader(
            dataset,
            batch,
            sampler=sampler,
            num_workers=args.num_workers,
            pin_memory=True,
            drop_last=False,
        )

        # generated images should match with n sample
        if n_sample % batch == 0:
            assert len(loader) == n_sample // batch
        else:
            assert len(loader) == n_sample // batch + 1

        total_latents = torch.Tensor().to(device)
        real_imgs = torch.Tensor().to(device)

        if args.mixing_type == "local_editing":
            if dataset_name == "afhq":
                masks = (
                    -2 * torch.ones(n_sample, args.size, args.size).to(device).float()
                )

                mix_type = list(range(n_sample))
                random.shuffle(mix_type)
                horizontal_mix = mix_type[: n_sample // 2]
                vertical_mix = mix_type[n_sample // 2 :]

                masks[horizontal_mix, :, args.size // 2 :] = 2
                masks[vertical_mix, args.size // 2 :, :] = 2
            else:
                masks = torch.Tensor().to(device).long()

    with torch.no_grad():
        if args.mixing_type == "random_generation":
            truncation = 0.7
            truncation_sample = 5000
            truncation_mean_latent = torch.Tensor().to(device)
            for _ in range(truncation_sample // batch):
                z = make_noise(batch, args.latent_channel_size, device)
                partial_mean_latent = model(z, mode="calculate_mean_stylemap")
                truncation_mean_latent = torch.cat(
                    [truncation_mean_latent, partial_mean_latent], dim=0
                )
            truncation_mean_latent = truncation_mean_latent.mean(0, keepdim=True)

            # refer to stylegan official repository: https://github.com/NVlabs/stylegan/blob/master/generate_figures.py
            cx, cy, cw, ch, rows, lods = 0, 0, 1024, 1024, 3, [0, 1, 2, 2, 3, 3]

            for seed in range(0, 4):
                torch.manual_seed(seed)
                png = f"{args.save_image_dir}/random_generation_{seed}.png"
                print(png)

                total_images_len = sum(rows * 2 ** lod for lod in lods)
                total_images = torch.Tensor()

                while total_images_len > 0:
                    num = batch if total_images_len > batch else total_images_len
                    z = make_noise(num, args.latent_channel_size, device)
                    total_images_len -= batch

                    images = model(
                        (z, truncation, truncation_mean_latent),
                        mode="random_generation",
                    )

                    images = images.permute(0, 2, 3, 1)
                    images = images.cpu()
                    total_images = torch.cat([total_images, images], dim=0)

                total_images = torch.clamp(total_images, min=-1.0, max=1.0)
                total_images = (total_images + 1) / 2 * 255
                total_images = total_images.numpy().astype(np.uint8)

                canvas = Image.new(
                    "RGB",
                    (sum(cw // 2 ** lod for lod in lods), ch * rows),
                    "white",
                )
                image_iter = iter(list(total_images))
                for col, lod in enumerate(lods):
                    for row in range(rows * 2 ** lod):
                        image = Image.fromarray(next(image_iter), "RGB")
                        # image = image.crop((cx, cy, cx + cw, cy + ch))
                        image = image.resize(
                            (cw // 2 ** lod, ch // 2 ** lod), Image.ANTIALIAS
                        )
                        canvas.paste(
                            image,
                            (
                                sum(cw // 2 ** lod for lod in lods[:col]),
                                row * ch // 2 ** lod,
                            ),
                        )
                canvas.save(png)

        elif args.mixing_type == "reconstruction":
            for i, real_img in enumerate(tqdm(loader, mininterval=1)):
                real_img = real_img.to(device)
                recon_image = model(real_img, "reconstruction")

                for i_b, (img_1, img_2) in enumerate(zip(real_img, recon_image)):
                    save_images(
                        [img_1, img_2],
                        [
                            f"{args.save_image_dir}/{i*batch+i_b}_real.png",
                            f"{args.save_image_dir}/{i*batch+i_b}_recon.png",
                        ],
                    )

        elif args.mixing_type == "transplantation":

            for kind in [
                "source_image",
                "source_reconstruction",
                "reference_image",
                "reference_reconstruction",
                "synthesized_image",
            ]:
                os.makedirs(os.path.join(args.save_image_dir, kind), exist_ok=True)

            # AFHQ
            transplantation_dataset = [
                (62, 271, [((4, 2), (3, 2), 2, 4), ((0, 1), (0, 1), 3, 2)])
            ]

            for index_src, index_ref, coordinates in transplantation_dataset:
                src_img = dataset[index_src].to(device)
                ref_img = dataset[index_ref].to(device)

                mixed_image, recon_img_src, recon_img_ref = model(
                    (src_img, ref_img, coordinates), mode="transplantation"
                )

                ratio = 256 // 8

                src_img = (src_img + 1) / 2
                ref_img = (ref_img + 1) / 2

                colors = [(0, 0, 255), (0, 255, 0), (0, 255, 0)]

                for color_i, (
                    (src_p_y, src_p_x),
                    (ref_p_y, ref_p_x),
                    height,
                    width,
                ) in enumerate(coordinates):
                    for i in range(2):
                        img = src_img if i == 0 else ref_img
                        img = img.cpu()
                        img = transforms.ToPILImage()(img)
                        img = np.asarray(img)
                        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
                        if i == 0:
                            img = cv2.rectangle(
                                img,
                                (src_p_x * ratio, src_p_y * ratio),
                                (
                                    (src_p_x + width) * ratio,
                                    (src_p_y + height) * ratio,
                                ),
                                colors[color_i],
                                2,
                            )
                        else:
                            img = cv2.rectangle(
                                img,
                                (ref_p_x * ratio, ref_p_y * ratio),
                                (
                                    (ref_p_x + width) * ratio,
                                    (ref_p_y + height) * ratio,
                                ),
                                colors[color_i],
                                2,
                            )
                        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                        img = transforms.ToTensor()(img)

                        if i == 0:
                            src_img = img
                        else:
                            ref_img = img

                save_images(
                    [mixed_image[0], recon_img_src[0], recon_img_ref[0]],
                    [
                        f"{args.save_image_dir}/synthesized_image/{index_src}_{index_ref}.png",
                        f"{args.save_image_dir}/source_reconstruction/{index_src}_{index_ref}.png",
                        f"{args.save_image_dir}/reference_reconstruction/{index_src}_{index_ref}.png",
                    ],
                )

                save_images(
                    [src_img, ref_img],
                    [
                        f"{args.save_image_dir}/source_image/{index_src}_{index_ref}.png",
                        f"{args.save_image_dir}/reference_image/{index_src}_{index_ref}.png",
                    ],
                    range=(0, 1),
                )

        else:
            for i, real_img in enumerate(tqdm(loader, mininterval=1)):
                if (args.mixing_type == "local_editing") and (
                    dataset_name == "celeba_hq"
                ):
                    real_img, mask = real_img
                    mask = mask.to(device)
                    masks = torch.cat([masks, mask], dim=0)
                real_img = real_img.to(device)

                latents = model(real_img, "projection")

                total_latents = torch.cat([total_latents, latents], dim=0)
                real_imgs = torch.cat([real_imgs, real_img], dim=0)

            if args.mixing_type == "stylemixing":
                with open(
                    "data/celeba_hq/local_editing/celeba_hq_test_GT_sorted_pair.pkl",
                    "rb",
                ) as f:
                    sorted_similarity = pickle.load(f)

                indices1 = []
                indices2 = []
                reference = [356, 358, 363, 400, 483]
                original = [7, 124, 162, 135, 136, 173]
                for i in original:
                    indices1.append(i)
                    indices2.append(reference)

            elif args.mixing_type == "w_interpolation":
                indices = list(range(len(total_latents)))

                if dataset_name == "celeba_hq":
                    n_sample = 29000
                elif dataset_name == "afhq":
                    n_sample = 15130
                elif dataset_name == "ffhq":
                    n_sample = 69000

                indices1 = random.choices(indices, k=n_sample)
                indices2 = random.choices(indices, k=n_sample)

            elif args.mixing_type == "local_editing":
                if dataset_name == "afhq":
                    # change it later
                    indices = list(range(len(total_latents)))
                    random.shuffle(indices)
                    indices1 = indices[: len(total_latents) // 2]
                    indices2 = indices[len(total_latents) // 2 :]

                else:
                    with open(
                        f"{mask_path_base}/celeba_hq_test_GT_sorted_pair.pkl",
                        "rb",
                    ) as f:
                        sorted_similarity = pickle.load(f)

                    indices1 = []
                    indices2 = []
                    for (i1, i2), _ in sorted_similarity[args.local_editing_part]:
                        indices1.append(i1)
                        indices2.append(i2)

            for loop_i, (index1, index2) in tqdm(
                enumerate(zip(indices1, indices2)), total=n_sample
            ):
                if args.mixing_type == "w_interpolation":
                    imgs = model(
                        (total_latents[index1], total_latents[index2]),
                        "w_interpolation",
                    )
                    assert len(imgs) == 1
                    save_image(
                        imgs[0],
                        f"{args.save_image_dir}/{loop_i}.png",
                    )
                elif args.mixing_type == "stylemixing":
                    n_rows = len(index2)
                    coarse_img, fine_img = model(
                        (
                            torch.stack([total_latents[index1] for _ in range(n_rows)]),
                            torch.stack([total_latents[i2] for i2 in index2]),
                        ),
                        "stylemixing",
                    )

                    save_images(
                        [coarse_img, fine_img],
                        [
                            f"{args.save_image_dir}/{index1}_coarse.png",
                            f"{args.save_image_dir}/{index1}_fine.png",
                        ],
                    )

                elif args.mixing_type == "local_editing":
                    src_img = real_imgs[index1]
                    ref_img = real_imgs[index2]

                    if dataset_name == "celeba_hq":
                        mask1_logit = masks[index1]
                        mask2_logit = masks[index2]

                        mask1 = -torch.ones(mask1_logit.shape).to(
                            device
                        )  # initialize with -1
                        mask2 = -torch.ones(mask2_logit.shape).to(
                            device
                        )  # initialize with -1

                        for label_i in parts_index[args.local_editing_part]:
                            mask1[(mask1_logit == label_i) == True] = 1
                            mask2[(mask2_logit == label_i) == True] = 1

                        mask = mask1 + mask2
                        mask = mask.float()
                    elif dataset_name == "afhq":
                        mask = masks[index1]

                    mixed_image, recon_img_src, recon_img_ref = model(
                        (total_latents[index1], total_latents[index2], mask),
                        "local_editing",
                    )

                    save_images(
                        [
                            mixed_image[0],
                            recon_img_src[0],
                            src_img,
                            ref_img,
                            recon_img_ref[0],
                        ],
                        [
                            f"{args.save_image_dir}/synthesized_image/{index1}.png",
                            f"{args.save_image_dir}/source_reconstruction/{index1}.png",
                            f"{args.save_image_dir}/source_image/{index1}.png",
                            f"{args.save_image_dir}/reference_image/{index1}.png",
                            f"{args.save_image_dir}/reference_reconstruction/{index1}.png",
                        ],
                    )

                    mask[mask < -1] = -1
                    mask[mask > -1] = 1

                    save_image(
                        mask,
                        f"{args.save_image_dir}/mask/{index1}.png",
                    )

                    if dataset_name == "celeba_hq":
                        save_images(
                            [mask1, mask2],
                            [
                                f"{args.save_image_dir}/mask_src/{index1}.png",
                                f"{args.save_image_dir}/mask_ref/{index1}.png",
                            ],
                        )


import os
import argparse
from setproctitle import setproctitle

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from tqdm import tqdm

import network
import misc_utils
from config import config
from dataset import train_dataset

def process(data_generater, num_gpus):
    images = []
    gt_boxes = []
    im_info = []
    for _ in range(num_gpus):
        data = data_generater.__next__()
        minibt_img, minibt_gt, minibt_info = \
                data['data'], data['boxes'], data['im_info']
        images.append(minibt_img)
        gt_boxes.append(minibt_gt)
        im_info.append(minibt_info)
    images = torch.Tensor(np.vstack(images)).cuda()
    gt_boxes = torch.Tensor(np.vstack(gt_boxes)).cuda()
    im_info = torch.Tensor(np.vstack(im_info)).cuda()
    return images, gt_boxes, im_info

def adjust_learning_rate(optimizer,epoch,lr):
    learning_rate = lr
    lr = learning_rate*(0.1**(epoch//10))
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr
    return lr

def train(params):
    total_nr_iters = config.train_base_iters
    batch_per_gpu = config.train_batch_per_gpu
    base_lr = config.base_lr
    line = 'network.base_lr.{}.train_iter.{}'.format(base_lr, total_nr_iters)
    print(line)
    # set model save path and log save path
    saveDir = config.model_dir
    misc_utils.ensure_dir(saveDir)
    fpath = os.path.join(config.output_dir, line+'.log')
    fid_log = open(fpath,'a')
    # set data input pipe
    program_name = config.program_name
    # check gpus
    torch.set_default_tensor_type('torch.FloatTensor')
    if not torch.cuda.is_available():
        print('No GPU exists!')
        return
    else:
        num_gpus = torch.cuda.device_count()
        train_iter = total_nr_iters//(num_gpus*batch_per_gpu)
        train_lr_decay = np.array(config.lr_decay)//(num_gpus*batch_per_gpu)
        train_dump_interval = config.model_dump_interval//(num_gpus*batch_per_gpu)
    train_lr = base_lr * num_gpus
    bt_size = num_gpus * batch_per_gpu
    line = 'Num of GPUs:{}, learning rate:{:.5f}, batch size:{},\
            train_iter:{}, decay_iter:{}, dump_interval:{}'.format(
            num_gpus,train_lr,bt_size,train_iter,train_lr_decay, train_dump_interval)
    print(line)
    print("Building netowrk.")
    net = network.Network()
    net.cuda()
    if params.resume_weights:
        model_file = os.path.join(saveDir, 'dump-{}.pth'.format(params.resume_weights))
        check_point = torch.load(model_file)
        net.load_state_dict(check_point['state_dict'])
    net = nn.DataParallel(net)
    # set the optimizer, use momentum and weight_decay
    optimizer = optim.SGD(net.parameters(), lr=train_lr, momentum=config.momentum, \
        weight_decay=config.weight_decay)
    # check if resume training
    training_data = train_dataset()
    net.train()
    if(params.progressbar):
        tqdm.monitor_interval = 0
        pbar = tqdm(total=train_iter, leave=False, ascii=True)
    dump_num = 1
    start_iter = 0
    if params.resume_weights:
        start_iter = int(params.resume_weights) * train_dump_interval
        if(start_iter >= train_lr_decay[0]):
            optimizer.param_groups[0]['lr'] = train_lr / 10
        if(start_iter >= train_lr_decay[1]):
            optimizer.param_groups[0]['lr'] = train_lr / 100
        dump_num = int(params.resume_weights) + 1
    for step in range(start_iter, train_iter):
        # warm up
        if step < config.warm_iter:
            alpha = step / config.warm_iter
            lr_new = 0.1 * train_lr + 0.9 * alpha * train_lr
            optimizer.param_groups[0]['lr'] = lr_new
        elif step == config.warm_iter:
            optimizer.param_groups[0]['lr'] = train_lr
        if step == train_lr_decay[0]:
            optimizer.param_groups[0]['lr'] = train_lr / 10
        elif step == train_lr_decay[1]:
            optimizer.param_groups[0]['lr'] = train_lr / 100
        # get training data
        images, gt_boxes, img_info = process(training_data, num_gpus)
        optimizer.zero_grad()
        # forwad
        outputs = net(images, img_info, gt_boxes)
        # collect the loss
        total_loss = sum([outputs[key].mean() for key in outputs.keys()])
        total_loss.backward()
        optimizer.step()
        if(params.progressbar):
            pbar.update(1)
        # stastic
        if step % config.log_dump_interval == 0:
            stastic_total_loss = total_loss.cpu().data.numpy()
            line = 'Iter {}: lr:{:.5f}, loss is {:.4f}.'.format(
                step, optimizer.param_groups[0]['lr'], stastic_total_loss)
            print(outputs)
            print(line)
            fid_log.write(line+'\n')
            fid_log.flush()
        # save the model
        if (step + 1)%train_dump_interval==0:
            fpath = os.path.join(saveDir,'dump-{}.pth'.format(dump_num))
            dump_num += 1
            model = dict(epoch = step,
                state_dict = net.module.state_dict(),
                optimizer = optimizer.state_dict())
            torch.save(model,fpath)
    if(params.progressbar):
        pbar.close()
    fid_log.close()

def run_train():
    setproctitle('train ' + os.path.split(os.path.realpath(__file__))[0])
    parser = argparse.ArgumentParser()
    parser.add_argument('--resume_weights', '-r', default=None,type=str)
    parser.add_argument('--progressbar', '-p', action='store_true', default=False)
    args = parser.parse_args()
    train(args)

if __name__ == '__main__':
    run_train()


# module imports
import random
import numpy as np

# universal colors
bg = (255, 255, 255)  # background color 'white'
ln = (0, 0, 0)        # line color 'black'
tg = (255, 0, 127)    # tongue color 'red' variant

# Random face color applying to all table values. Starting with 50 range value produces brighter colors
rf = random.randint(50, 255)
gf = random.randint(50, 255)
bf = random.randint(50, 255)
face_random = [rf, gf, bf]
ce = face_random

# Random body color based on RGB value range
rb = random.randint(0, 255)
gb = random.randint(0, 255)
bb = random.randint(0, 255)
body_random = [rb, gb, bb]
dy = body_random

# Define random eye colors if rolled
re1 = random.randint(100, 255)
ge1 = random.randint(100, 255)
be1 = random.randint(100, 255)
eye1_random = [re1, ge1, be1]

re2 = random.randint(50, 255)
ge2 = random.randint(50, 255)
be2 = random.randint(50, 255)
eye2_random = [re2, ge2, be2]

re3 = random.randint(0, 255)
ge3 = random.randint(0, 255)
be3 = random.randint(0, 255)
eye3_random = [re3, ge3, be3]

# define random Snake Skin colors if rolled
rss = random.randint(0, 255)
gss = random.randint(0, 255)
bss = random.randint(0, 255)
snakeskin_random = [rss, gss, bss]

# chance for Hypnotize roll
for h in range(100):
        h = random.randint(0, 100)

        # Random number greater or equal to 95 generates Hypnotize
        if h >= 95:
                ey1 = eye1_random
                ey2 = eye2_random
                ey3 = eye3_random

        else:  # If unlucky, Reassign randomized eyes to RGB White
                ey1 = (255, 255, 255)
                ey2 = (255, 255, 255)
                ey3 = (255, 255, 255)

# chance for Snake Skin roll
for s in range(100):
        s = random.randint(0, 100)

        # Random number greater or equal to 95 generates Snake Skin
        if s >= 95:
                bb = snakeskin_random
        else:   # if unlucky bb values stay random body values defined above
                bb = dy

# chance for Fang roll
for f in range(100):
        f = random.randint(0, 100)

        # random number greater or equal to 92 generates Fangs
        if f >= 92:
                fa = (255, 255, 255)  # Fangs RGB value 'White'
        else: # if unlucky fa values stay random body values defined above
                fa = dy

# value array tables where Pythons are drawn with pixel value tied to "x" input

# garter type array table values
array_1 = np.array([
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg], # 1
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg], # 2
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg], # 3
        [bg, bg, bg, bg, bg, bg, bg, ln, ln, ln, ln, bg, bg, ln, ln, ln, bg, bg, bg, bg, bg, bg, bg, bg], # 4
        [bg, bg, bg, bg, bg, bg, ln, ce, ey1, ey2, ey1, ln, ln, ey3, ey2, ey3, ln, bg, bg, bg, bg, bg, bg, bg], # 5
        [bg, bg, bg, bg, bg, ln, ce, ce, ey2, ey3, ln, ce, ce, ey2, ey1, ln, ce, ln, bg, bg, bg, bg, bg, bg], # 6
        [bg, bg, bg, bg, bg, ln, dy, ce, ey3, ey2, ln, ce, ce, ey1, ey3, ln, ce, ce, ln, bg, bg, bg, bg, bg], # 7
        [bg, bg, bg, bg, bg, ln, dy, ce, ce, ce, ce, ce, ce, ce, ce, ce, ce, ce, ln, bg, bg, bg, bg, bg], # 8
        [bg, bg, bg, bg, bg, ln, dy, ce, ce, ce, ce, ce, ce, ce, ce, ce, ce, dy, ln, bg, bg, bg, bg, bg], # 9
        [bg, bg, bg, bg, bg, ln, dy, dy, ce, ce, ce, ce, ce, ce, ce, ce, ce, dy, ln, bg, bg, bg, bg, bg], # 10
        [bg, bg, bg, bg, bg, bg, ln, ln, dy, ce, ce, ce, ce, ce, ce, dy, ln, ln, bg, bg, bg, bg, bg, bg], # 11
        [bg, bg, bg, bg, bg, bg, bg, ln, dy, ce, ce, ce, ce, ce, ce, dy, ln, bg, bg, bg, bg, bg, bg, bg], # 12
        [bg, bg, bg, bg, bg, bg, bg, ln, dy, dy, ce, ln, dy, ln, ce, dy, ln, bg, bg, bg, bg, bg, bg, bg], # 13
        [bg, bg, bg, bg, bg, bg, bg, bg, ln, dy, dy, dy, dy, dy, dy, ln, bg, bg, bg, bg, bg, bg, bg, bg], # 14
        [bg, bg, bg, bg, bg, bg, bg, bg, ln, ln, ln, ln, tg, ln, ln, ln, bg, bg, bg, bg, bg, bg, bg, bg], # 15
        [bg, bg, bg, bg, bg, bg, bg, bg, ln, ce, fa, dy, tg, fa, ce, ln, bg, bg, bg, bg, bg, bg, bg, bg], # 16
        [bg, bg, bg, bg, bg, bg, bg, bg, ln, ce, dy, tg, tg, dy, ce, ln, bg, bg, bg, bg, bg, bg, bg, bg], # 17
        [bg, bg, bg, bg, bg, bg, bg, bg, ln, ce, bb, tg, bb, bb, ce, ln, bg, bg, bg, bg, bg, bg, bg, bg], # 18
        [bg, bg, bg, bg, bg, bg, bg, bg, ln, ce, dy, dy, dy, dy, ce, ln, bg, bg, bg, bg, bg, bg, bg, bg], # 19
        [bg, bg, bg, bg, bg, bg, bg, bg, ln, ce, bb, bb, bb, bb, ce, ln, bg, bg, bg, bg, bg, bg, bg, bg], # 20
        [bg, bg, bg, bg, bg, bg, bg, bg, ln, ce, dy, dy, dy, dy, ce, ln, bg, bg, bg, bg, bg, bg, bg, bg], # 21
        [bg, bg, bg, bg, bg, bg, bg, bg, ln, ce, bb, bb, bb, bb, ce, ln, bg, bg, bg, bg, bg, bg, bg, bg], # 22
        [bg, bg, bg, bg, bg, bg, bg, bg, ln, ce, dy, dy, dy, dy, ce, ln, bg, bg, bg, bg, bg, bg, bg, bg], # 23
        [bg, bg, bg, bg, bg, bg, bg, bg, ln, ce, bb, bb, bb, bb, ce, ln, bg, bg, bg, bg, bg, bg, bg, bg]])  # 24

# rattler type array table values
array_2 =np.array([
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg], # 1
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg], # 2
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg], # 3
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg], # 4
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg], # 5
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg], # 6
        [bg, bg, bg, bg, bg, ln, ln, ln, ln, ln, ln, ln, ln, ln, ln, ln, ln, ln, bg, bg, bg, bg, bg, bg], # 7
        [bg, bg, bg, bg, ln, ce, ce, ce, ce, ce, ce, ce, ce, ce, ce, ce, ce, ce, ln, bg, bg, bg, bg, bg], # 8
        [bg, bg, bg, ln, ce, ce, ce, ce, ce, ce, ce, ce, ce, ce, ce, ce, ce, ce, ce, ln, bg, bg, bg, bg], # 9
        [bg, bg, bg, ln, ce, ce, bb, dy, bb, ln, ln, ln, ce, ce, ce, ce, ce, ce, ce, ln, bg, bg, bg, bg], # 10
        [bg, bg, bg, ln, ce, bb, dy, bb, ln, ce, ce, ce, ce, ce, ce, ce, ce, ce, ce, ln, bg, bg, bg, bg], # 11
        [bg, bg, bg, ln, ln, ln, ln, ln, ln, ln, ln, ln, ln, ce, ce, ce, ce, ce, ce, ln, bg, bg, bg, bg], # 12
        [bg, bg, ln, ce, ce, ce, ce, ce, ce, ce, ce, ce, ce, ln, ce, ce, ln, ln, ln, bg, bg, bg, bg, bg], # 13
        [bg, bg, ln, bb, ce, ey1, ey2, ce, ce, ce, ey3, ey1, ce, ln, ln, ln, bg, bg, bg, bg, bg, bg, bg, bg], # 14
        [bg, bg, ln, dy, ce, ey3, ln, ce, ce, ce, ey2, ln, ce, ln, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg], # 15
        [bg, bg, ln, bb, ce, ce, ce, ce, ce, ce, ce, ce, ce, ln, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg], # 16
        [bg, bg, bg, ln, bb, dy, ce, ce, ce, bb, ce, ce, bb, ln, bg, bg, bg, bg, ln, bg, bg, bg, bg, bg], # 17
        [bg, bg, bg, ln, ln, bb, dy, bb, ln, dy, ln, bb, ln, bg, bg, bg, bg, ln, bb, ln, bg, bg, bg, bg], # 18
        [bg, bg, bg, ln, ce, ln, ln, ln, ln, ln, ln, ln, bg, bg, bg, bg, bg, ln, dy, ln, bg, bg, bg, bg], # 19
        [bg, bg, bg, ln, ce, bb, fa, bb, bb, fa, ln, bg, bg, bg, bg, bg, bg, ln, tg, ln, bg, bg, bg, bg], # 20
        [bg, bg, bg, ln, ce, fa, dy, bb, fa, bb, ln, bg, bg, bg, bg, bg, bg, ln, dy, ln, bg, bg, bg, bg], # 21
        [bg, bg, bg, ln, ce, dy, bb, bb, dy, dy, bb, ln, bg, bg, bg, bg, ln, dy, bb, ln, bg, bg, bg, bg], # 22
        [bg, bg, bg, ln, ce, bb, dy, dy, bb, dy, dy, ln, bg, bg, bg, bg, ln, tg, tg, ln, bg, bg, bg, bg], # 23
        [bg, bg, bg, ln, ce, dy, bb, dy, dy, bb, bb, ln, bg, bg, bg, ln, bb, dy, ln, bg, bg, bg, bg, bg]]) # 24

# cobra type array table values
array_3 = np.array([
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg], # 1
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg], # 2
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg], # 3
        [bg, bg, bg, bg, bg, bg, bg, bg, ln, ln, ln, ln, ln, ln, ln, ln, ln, bg, bg, bg, bg, bg, bg, bg], # 4
        [bg, bg, bg, bg, bg, bg, bg, ln, ce, ce, ce, ce, ce, ce, ce, ce, ce, ln, bg, bg, bg, bg, bg, bg], # 5
        [bg, bg, bg, bg, bg, bg, bg, ln, ce, ey1, ln, ce, ce, ey2, ln, ce, ce, ce, ln, bg, bg, bg, bg, bg], # 6
        [bg, bg, bg, bg, bg, bg, ln, ln, ce, ey3, ln, ce, ce, ey1, ln, ce, ce, ce, ce, ln, bg, bg, bg, bg], # 7
        [bg, bg, bg, bg, bg, ln, ce, ln, ce, ce, ce, ce, ce, ce, ce, ce, ce, ln, ce, ce, ln, bg, bg, bg], # 8
        [bg, bg, bg, bg, ln, ce, bb, ln, ce, ce, ce, ce, ce, ce, ce, ce, ce, ln, bb, ce, ln, bg, bg, bg], # 9
        [bg, bg, bg, bg, ln, ce, dy, ln, ce, ce, ln, ce, ln, ce, ce, ce, ln, dy, dy, ce, ln, bg, bg, bg], # 10
        [bg, bg, bg, bg, ln, ce, dy, ln, ln, ce, ce, ce, ce, ce, ce, ln, dy, dy, dy, ce, ln, bg, bg, bg], # 11
        [bg, bg, bg, bg, ln, ce, bb, bb, ln, ln, ln, tg, ln, ln, ln, bb, bb, bb, bb, ce, ln, bg, bg, bg], # 12
        [bg, bg, bg, bg, ln, ce, dy, dy, dy, fa, dy, tg, dy, dy, fa, dy, dy, dy, dy, ce, ln, bg, bg, bg], # 13
        [bg, bg, bg, bg, ln, ce, dy, dy, dy, fa, dy, tg, dy, dy, fa, dy, dy, dy, dy, ce, ln, bg, bg, bg], # 14
        [bg, bg, bg, bg, ln, ce, bb, bb, bb, bb, tg, tg, tg, bb, bb, bb, bb, bb, bb, ce, ln, bg, bg, bg], # 15
        [bg, bg, bg, bg, bg, ln, ce, dy, dy, dy, tg, dy, tg, dy, dy, dy, dy, dy, ce, ce, ln, bg, bg, bg], # 16
        [bg, bg, bg, bg, bg, ln, ce, dy, dy, dy, dy, dy, dy, dy, dy, dy, dy, dy, ce, ln, ln, bg, bg, bg], # 17
        [bg, bg, bg, bg, bg, bg, ln, ce, bb, bb, bb, bb, bb, bb, bb, bb, dy, ce, ce, ln, bg, bg, bg, bg], # 18
        [bg, bg, bg, bg, bg, bg, bg, ln, dy, dy, dy, dy, dy, dy, dy, dy, ce, ce, ln, bg, bg, bg, bg, bg], # 19
        [bg, bg, bg, bg, bg, bg, bg, ln, dy, dy, dy, dy, dy, dy, dy, ce, ce, ln, bg, bg, bg, bg, bg, bg], # 20
        [bg, bg, bg, bg, bg, bg, bg, ln, bb, bb, bb, bb, bb, dy, ce, ln, ln, bg, bg, bg, bg, bg, bg, bg], # 21
        [bg, bg, bg, bg, bg, bg, bg, ln, dy, dy, dy, dy, dy, dy, ce, ln, bg, bg, bg, bg, bg, bg, bg, bg], # 22
        [bg, bg, bg, bg, bg, bg, bg, ln, dy, dy, dy, dy, dy, dy, ce, ln, bg, bg, bg, bg, bg, bg, bg, bg], # 23
        [bg, bg, bg, bg, bg, bg, bg, ln, bb, bb, bb, bb, dy, dy, ce, ln, bg, bg, bg, bg, bg, bg, bg, bg]])  # 24

# mamba type array table values
array_4 = np.array([
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg], # 1
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg], # 2
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, ln, ln, ln, ln, bg, bg, bg, bg, bg, bg, bg, bg], # 3
        [bg, bg, bg, bg, bg, bg, ln, ln, ln, ln, ln, ln, ce, ln, ce, ce, ln, bg, bg, bg, bg, bg, bg, bg], # 4
        [bg, bg, bg, bg, bg, ln, ce, ce, ey1, ln, ey3, ce, ce, ce, ce, ce, ce, ln, bg, bg, bg, bg, bg, bg], # 5
        [bg, bg, bg, bg, bg, ln, ce, ce, ce, ce, ce, ce, ce, fa, dy, dy, fa, ln, bg, bg, bg, bg, bg, bg], # 6
        [bg, bg, bg, bg, ln, ce, ce, ce, ce, ce, ce, dy, dy, fa, dy, dy, fa, ln, bg, bg, bg, bg, bg, bg], # 7
        [bg, bg, bg, bg, ln, ce, ce, ce, ce, ce, dy, dy, fa, dy, dy, fa, ln, bg, bg, bg, bg, bg, bg, bg], # 8
        [bg, bg, bg, bg, ln, ce, ce, ce, ln, dy, dy, dy, dy, dy, dy, ln, bg, bg, bg, bg, bg, bg, bg, bg], # 9
        [bg, bg, bg, bg, ln, ce, ce, ce, ln, dy, dy, dy, dy, dy, dy, ln, bg, bg, bg, bg, bg, bg, bg, bg], # 10
        [bg, bg, bg, bg, ln, ce, ce, ce, ce, ln, dy, ln, ln, dy, ln, bg, bg, bg, bg, bg, bg, bg, bg, bg], # 11
        [bg, bg, bg, bg, ln, ce, ce, ce, ce, ln, dy, dy, ln, dy, dy, ln, bg, bg, bg, bg, bg, bg, bg, bg], # 12
        [bg, bg, bg, bg, ln, ce, ce, ce, ce, dy, ln, dy, dy, ln, ln, dy, ln, bg, bg, bg, bg, bg, bg, bg], # 13
        [bg, bg, bg, bg, ln, ce, ce, ce, dy, dy, ln, ln, dy, ln, dy, dy, ln, bg, bg, bg, bg, bg, bg, bg], # 14
        [bg, bg, bg, bg, ln, ce, ce, ce, bb, bb, ln, bg, ln, ln, ln, ln, bg, bg, bg, bg, bg, bg, bg, bg], # 15
        [bg, bg, bg, bg, ln, ce, ce, ce, bb, bb, ln, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg], # 16
        [bg, bg, bg, bg, ln, ce, ce, ce, dy, dy, ln, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg], # 17
        [bg, bg, bg, bg, ln, ce, ce, ce, dy, dy, ln, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg], # 18
        [bg, bg, bg, bg, ln, ce, ce, ce, bb, bb, ln, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg], # 19
        [bg, bg, bg, bg, ln, ce, ce, ce, bb, bb, ln, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg], # 20
        [bg, bg, bg, bg, ln, ce, ce, ce, dy, dy, ln, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg], # 21
        [bg, bg, bg, bg, ln, ce, ce, ce, dy, dy, ln, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg], # 22
        [bg, bg, bg, bg, ln, ce, ce, ce, bb, bb, ln, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg], # 23
        [bg, bg, bg, bg, ln, ce, ce, ce, bb, bb, ln, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg]]) # 24

# this file is pulled into the main file where the body type is rolled and image is generated