import torch.nn.functional as F
import torch.nn as nn
import torch
import math
import torch
import numpy as np
from opacus.layers.dp_rnn import DPGRUCell
from dataset import TrajectoryDataset
from my_utils import construct_default_quadtree

class BaseReferenceGenerator(nn.Module):
    """
    reference includes trajectory type and start_location
    this generates a trajectory that starts with start location and follows the trajectory type
    e.g., x, 1, 0 -> x, ?, x
    """

    def post_process(self, output, sampled, reference):
        """
        reference is a tuple of ints
        the length of the reference is the length of the trajectory
        the first element is the start location
        the rest of the elements are references to a first appearance of the location in the trajectory
        i.e., (x, 1, 0, 1) induces (x, y, x, y) where x and y are locations
        """
        # if the output is a list, this is the distributions for all layers
        output = output[-1] if type(output) == list else output
        assert len(reference) == output.shape[0], "the length of the reference should be {}, but it is {}".format(output.shape[0], len(reference))

        # remove reference complementation for fair comparison
        if True:
            # test!
            pointer = len(sampled)
            if pointer == 1:
                if reference[0][0] != -1:
                    locations = torch.tensor([v[0] for v in reference]).view(-1, 1).to(output.device)
            else:
                location_probs = torch.exp(output).view(output.shape[0], -1)
                locations = location_probs.multinomial(1)
            return locations

        log_location_probs = output

        # pointer = len(sampled)
        # if pointer == 1:
        #     if reference[0][0] != -1:
        #         locations = torch.tensor([v[0] for v in reference]).view(-1, 1).to(output.device)
        #     elif torch.allclose(reference[:, 0], -1*torch.ones_like(reference[:, 0])):
        #         location_probs = torch.exp(log_location_probs).view(log_location_probs.shape[0], -1)
        #         locations = location_probs.multinomial(1)
        #     else:
        #         raise NotImplementedError("the first element of the reference should be -1, but it is {}".format(reference[0][0]))
        # else:
        #     passed_locations = torch.concatenate([v for v in sampled[1:]], dim=1).view(-1, len(sampled)-1)
        #     log_location_probs = self.remove_location(log_location_probs, passed_locations)
        #     location_probs = torch.exp(log_location_probs).view(log_location_probs.shape[0], -1)
        #     # if all locations are removed, then we just sample from the first location
        #     indice = location_probs.sum(dim=-1) == 0
        #     location_probs[indice,0] = 1
        #     # sample
        #     locations = location_probs.multinomial(1)
        #     # replace by the reference
        #     locations = torch.concat([passed_locations, locations], dim=-1)[range(len(reference)), reference[:, pointer-1]].view(-1, 1)

        return locations

    def step(self, input):
        return self(input)

    def make_initial_input(self, reference):
        batch_size = len(reference)
        seq_len = max([len(ref) for ref in reference])
        outputs = torch.tensor([self.start_idx for _ in range(batch_size)]).view(batch_size, 1)
        return outputs, seq_len

    def concat(self, sampled, samples, reference):
        if samples == None:
            samples = []
        sampled = torch.concat(sampled, dim=1)
        samples.extend(sampled.cpu().detach().numpy().tolist())
        return samples

    def make_sample(self, references, batch_size):
        '''
        make n_samples of trajectories usign minibatch of batch_size
        kwargs are used as the auxiliary information to post-process output
        '''
        n_samples = len(references)

        samples = None
        for i in range(int(n_samples / batch_size)):
            reference = references[i*batch_size:(i+1)*batch_size]
            outputs, seq_len = self.make_initial_input(reference)
            reference_input = torch.tensor([list(ref) + [i for i in range(len(ref), seq_len)] for ref in reference])
            sampled = [outputs]
            for j in range(seq_len):
                outputs = self.step(outputs)
                # the main function of post_process is sample a location from the corresponding distribution
                outputs = self.post_process(outputs, sampled, reference_input)
                sampled.append(outputs)
            sampled = sampled[1:]

            samples = self.concat(sampled, samples, reference)
        return samples

    def remove_location(self, location, remove_locationss, log_prob=True):
        # remove the locations that are in remove_locations
        if remove_locationss is not None:
            assert location.shape[0] == len(remove_locationss)
            for i, remove_locations in enumerate(remove_locationss):
                if log_prob:
                    location[i, ..., remove_locations] = -float("inf")
                else:
                    location[i, ..., remove_locations] = 0
        if not log_prob:
            location = F.normalize(location, p=1, dim=-1)
        return location

class Markov1Generator(BaseReferenceGenerator):
    def __init__(self, transition_matrix, state_to_class):
        """
        transition_matrix: class * n_locations
        """
        n_classes = transition_matrix.shape[0]
        n_locations = len(state_to_class)
        assert transition_matrix.shape[1] == n_locations, "the number of locations should be {}, but it is {}".format(n_locations, transition_matrix.shape[1])
        start_prob = torch.zeros_like(transition_matrix[1]).view(1, -1)
        start_prob[0][0] = 1
        self.transition_matrix = torch.concat([transition_matrix, start_prob], dim=0)
        self.state_to_class = state_to_class
        self.start_idx = n_locations
        self.state_to_class[self.start_idx] = n_classes

    def __call__(self, input):
        """
        input is supposed to be batched location: batch_size * 1
        start_idx is going to be 0
        """
        classes = torch.tensor([self.state_to_class[state.item()] for state in input])
        output = self.transition_matrix[classes]
        return torch.log(output)

    def train(self):
        pass

    def eval(self):
        pass

class BaseTimeReferenceGenerator(BaseReferenceGenerator):
    """
    reference includes trajectory type and start_location
    this generates a trajectory that starts with start location and follows the trajectory type
    e.g., x, 1, 0 -> x, ?, x
    """

    def post_process(self, output, sampled, reference):
        """
        output is a tuple of (location, time)
        location is post processed by super().post_process
        time is sampled from the time distribution
        time starts with 1 and a next time should be larger than or equal to the previous time
        """

        log_location_probs, log_time_probs = output
        locations = super().post_process(log_location_probs, [v[0] for v in sampled], reference)
        if len(sampled) == 1:
            times = torch.zeros_like(locations)
        else:
            # choose a time that is larger than the previous time
            previous_times = sampled[-1][1]
            passed_times = [list(range(time-1)) for time in previous_times]
            log_time_probs = self.remove_location(log_time_probs, passed_times)
            times = torch.exp(log_time_probs).multinomial(1)+1

        return [locations, times]

    def concat(self, sampled, samples, reference):
        if samples == None:
            samples = [[], []]
        locations = torch.concat([v[0] for v in sampled], dim=1).cpu().detach().numpy().tolist()
        times = torch.concat([v[1] for v in sampled], dim=1).cpu().detach().numpy().tolist()

        # remove the outside of the format
        for i in range(len(locations)):
            samples[0].append(locations[i][:len(reference[i])])
            samples[1].append(times[i][:len(reference[i])])

        return samples

class GRUNet(BaseTimeReferenceGenerator):

    def __init__(self, n_locations, location_embedding_dim, time_dim, traj_type_dim, hidden_dim, reference_to_label):
        super(GRUNet, self).__init__()
        self.traj_type_embedding = nn.Embedding(traj_type_dim, hidden_dim)
        self.reference_to_label = reference_to_label
        self.n_locations = n_locations
        self.time_dim = time_dim
        self.hidden_dim = hidden_dim
        self.gru = DPGRUCell(location_embedding_dim+time_dim, hidden_dim, True)
        self.fc = nn.Linear(hidden_dim, n_locations+time_dim-1)
        self.location_embedding = nn.Embedding(TrajectoryDataset.vocab_size(n_locations), location_embedding_dim)
        self.start_idx = TrajectoryDataset.start_idx(n_locations)
        self.relu = nn.ReLU()

    def make_initial_input(self, reference):
        batch_size = len(reference)
        seq_len = max([len(ref) for ref in reference])
        locations = torch.tensor([self.start_idx for _ in range(batch_size)]).to(next(self.parameters()).device).view(batch_size, 1)
        times = torch.tensor([0 for _ in range(batch_size)]).to(next(self.parameters()).device).view(batch_size, 1)
        states = self.init_hidden(reference)
        return [locations, times, states], seq_len

    def input_gru(self, x, references):
        state = self.init_hidden(references)
        locations = x[0]
        times = x[1]

        seq_len = locations.shape[1]

        states = []
        for i in range(seq_len):
            input_locations = locations[:, i]
            input_times = times[:, i]
            state = self.encode(input_locations, input_times, state)
            states.append(state)

        out = torch.stack(states, dim=1)

        return out

    # from the input, predict the next location probs and next time probs
    # input sequence fo (embed location, embed time): batch_size * seq_len * (embed_size+time_dim)
    # output location, time: seq_len * batch_size * (location_dim, time_dim)
    def forward(self, x, references):
        # embedding prefix
        out = self.input_gru(x, references)

        # decoding to probability of next location and next time
        location, time = self.decode(out)
        return location, time

    def init_hidden(self, references):
        # labels = torch.tensor([self.reference_to_label(reference) for reference in references]).to(next(self.parameters()).device)
        # print(labels)
        labels = torch.zeros(len(references), dtype=int).to(next(self.parameters()).device)
        hidden = self.traj_type_embedding(labels).view(-1, self.hidden_dim)
        return hidden

    def post_process(self, output, sampled, reference):

        log_location_probs, log_time_probs, states = output
        locations, times = super().post_process([log_location_probs, log_time_probs], [[v[0], v[1]] for v in sampled], reference)

        return [locations, times, states]

    def step(self, input):
        # input: [locations, times, states]
        next_state = self.encode(*input)
        location_probs, time_probs = self.decode(next_state)
        if type(location_probs) == list:
            location_probs = location_probs[-1]

        return [location_probs, time_probs, next_state]

    def encode(self, location, time, state):
        x = self.location_embedding(location).view(location.shape[0], -1)
        # one hot encode time
        times = F.one_hot(time, num_classes=self.time_dim).long().view(time.shape[0], -1)
        # concat time information
        x = torch.cat([x, times], dim=-1)

        state = self.gru(x, state)
        return state

    def decode(self, encoded):
        out = self.fc(self.relu(encoded))
        # split the last dimension into location and time
        log_location_probs, log_time_probs = torch.split(out, [self.n_locations, self.time_dim-1], dim=-1)
        log_location_probs = F.log_softmax(log_location_probs, dim=-1)
        log_time_probs = F.log_softmax(log_time_probs, dim=-1)

        return log_location_probs, log_time_probs

class MetaGRUNet(GRUNet):
    def __init__(self, meta_network, n_locations, location_embedding_dim, time_dim, traj_type_dim, hidden_dim, reference_to_label):
        super(MetaGRUNet, self).__init__(n_locations, location_embedding_dim, time_dim, traj_type_dim, hidden_dim, reference_to_label)

        self.meta_net = meta_network
        del self.location_embedding
        if hasattr(self.meta_net, "location_embedding"):
            # assert location_embedding_dim == self.meta_net.memory_dim, "location_embedding_dim should be {}, but it is {}".format(self.meta_net.memory_dim, location_embedding_dim)
            self.location_embedding = self.meta_net.location_embedding
        else:
            self.location_embedding = nn.Embedding(TrajectoryDataset.vocab_size(n_locations), location_embedding_dim)

        self.fc = None
        self.fc_location = nn.Linear(hidden_dim, self.meta_net.input_dim)
        self.fc_time = nn.Linear(hidden_dim, time_dim-1)

    def decode(self, embedding):
        out = self.fc_location(embedding)
        location = self.meta_net(out)
        # when meta net returns a list, it is multi-resolution task learning
        time = F.log_softmax(self.fc_time(embedding), dim=-1)
        return location, time

def compute_loss_meta_gru_net(target_locations, target_times, output_locations, output_times, coef_location, coef_time):
    # list type means multi-resolution task learning
    if type(target_locations) != list:
        output_locations = [output_locations[-1]] if type(output_locations) == list else [output_locations]
        target_locations = [target_locations]

    n_locations = output_locations[-1].shape[-1]
    loss = []
    for i in range(len(target_locations)):
        location_dim = output_locations[i].shape[-1]
        output_locations[i] = output_locations[i].view(-1, location_dim)
        target_locations[i] = target_locations[i].view(-1)
        # coef = (i+1)/len(target_locations)
        coef = 1
        loss.append(coef*F.nll_loss(output_locations[i], target_locations[i], ignore_index=TrajectoryDataset.ignore_idx(n_locations)) * coef_location)
    loss.append(F.nll_loss(output_times.view(-1, output_times.shape[-1]), (target_times).view(-1)) * coef_time)
    return loss

def compute_loss_gru_meta_gru_net(target_paths, target_times, output_locations, output_times, coef_location, coef_time):
    output_locations = output_locations.view(-1, 4)
    target_paths = target_paths.view(-1)
    loss_location = torch.nn.functional.nll_loss(output_locations, target_paths, ignore_index=4, reduction='mean') * coef_location
    loss_time = torch.nn.functional.nll_loss(output_times.view(-1, output_times.shape[-1]), target_times.view(-1)) * coef_time
    return loss_location, loss_time

class MetaNetwork(nn.Module):
    '''
    this class is a meta network, does not have any memory funtion
    that is, the hidden state is directory converted to log distribution on locations by MLP
    '''

    def __init__(self, memory_hidden_dim, hidden_dim, output_dim, n_classes, activate):
        '''
        output_dim := location_dim
        '''
        super(MetaNetwork, self).__init__()
        if activate == "relu":
            self.activate = nn.ReLU()
        elif activate == "leaky_relu":
            self.activate = nn.LeakyReLU()
        elif activate == "sigmoid":
            self.activate = nn.Sigmoid()
        else:
            raise NotImplementedError("activate should be relu or sigmoid, but it is {}".format(activate))

        self.input_dim = memory_hidden_dim
        self.query_to_scores = nn.Sequential(nn.Linear(hidden_dim, hidden_dim), self.activate, nn.Linear(hidden_dim, output_dim))
        self.class_to_query = nn.Sequential(nn.Linear(n_classes, hidden_dim), self.activate, nn.Linear(hidden_dim, hidden_dim))
        self.hidden_to_query_ = nn.Sequential(nn.Linear(memory_hidden_dim, hidden_dim), self.activate, nn.Linear(hidden_dim, hidden_dim))
        self.n_classes = n_classes

    def forward(self, hidden):
        # hidden: batch_size * (seq_len) * n_classes
        hidden = hidden.view(hidden.shape[0], -1, hidden.shape[-1])

        query = self.hidden_to_query(hidden)
        # print(hidden.shape, query.shape)
        scores = self.compute_scores(query)
        return F.log_softmax(scores, dim=-1)

    def hidden_to_query(self, hidden):
       # for pre_training
        # if hidden.shape[-1] == self.n_classes:
        if hasattr(self, "class_to_query"):
            # hidden: batch_size * seq_len * n_classes
            query = self.class_to_query(hidden)
        else:
            query = self.hidden_to_query_(hidden)
        # output hidden: batch_size * seq_len * memory_dim
        return query

    def compute_scores(self, query):
        scores = self.query_to_scores(query)
        return scores

    def remove_class_to_query(self):
        del self.class_to_query
        # self.class_to_query.requires_grad_(False)

class BaseQuadTreeNetwork(nn.Module):
    def __init__(self, n_locations, memory_dim, hidden_dim, n_classes, activate, is_consistent=True):
        super(BaseQuadTreeNetwork, self).__init__()

        if activate == "relu":
            self.activate = nn.ReLU()
        elif activate == "leaky_relu":
            self.activate = nn.LeakyReLU()
        elif activate == "sigmoid":
            self.activate = nn.Sigmoid()
        else:
            raise NotImplementedError("activate should be relu or sigmoid, but it is {}".format(activate))

        self.tree = construct_default_quadtree(int(np.sqrt(n_locations))-2)
        self.tree.make_self_complete()
        self.memory_dim = memory_dim
        self.n_locations = len(self.tree.get_leafs())
        self.n_classes = n_classes
        self.class_to_query = nn.Linear(n_classes, self.memory_dim)
        # self.hidden_to_query_ = nn.Sequential(nn.Linear(hidden_dim, hidden_dim), self.activate, nn.Linear(hidden_dim, self.memory_dim))
        self.hidden_to_query_ = nn.Linear(hidden_dim, self.memory_dim)
        self.is_consistent = is_consistent
        self.pre_training = True

    def forward(self, hidden):

        if len(hidden.shape) == 2:
            hidden = hidden.view(hidden.shape[0], 1, hidden.shape[1])

        query = self.hidden_to_query(hidden)
        # print(hidden.shape, query.shape)
        scores = self.compute_scores(query)
        distributions_for_all_depths = self.to_location_distribution(scores, target_depth=0)
        return distributions_for_all_depths

    def compute_scores(self, query):
        '''
        scores: batch_size * (seq_len) * (4^depth+4^depth-1+...+4)
        '''
        original_shape = query.shape
        keys = self.make_keys(query.shape)

        # matmal key and field
        query = query.view(original_shape[0] * original_shape[1], -1, self.memory_dim)
        keys = keys.view(original_shape[0] * original_shape[1], -1, self.memory_dim)
        scores = torch.bmm(keys, query.transpose(-2,-1)).view(*original_shape[:-1], -1, 4)
        return scores

    def hidden_to_query(self, hidden):
       # for pre_training
        # if hidden.shape[-1] == self.n_classes:
        if hasattr(self, "class_to_query"):
            query = self.class_to_query(hidden)
        else:
            query = self.hidden_to_query_(hidden)
        return query

    def make_states(self, key):
        '''
        convert key to keys of the nodes in the tree
        if target is given, keys are the nodes in the path that is determined by the target
        for each key, target can specify a path (list of inputs with length depth)
        therefore, the shape of the target should be the same as the shape of the input except the last dimension
        that is, when target = None -> output: batch_size * (seq_len) * (4^depth+4^depth-1+...+4) * memory_dim
        when target != None -> output: batch_size * (seq_len) * depth * memory_dim
        '''
        pass

    def make_keys(self, shape):
        '''
        convert state to key by self.state_to_key (MLP)
        shaep should be [batch_size, seq_len, memory_dim]
        output: batch_size * seq_len * (4^depth+4^depth-1+...+4) * memory_dim
        the key is aligned with the order of the node_id in the tree
        '''
        states = self.make_states(shape)

        cursor = 0
        keys = []
        for i in range(1,1+self.tree.max_depth):
            keys_i = self.state_to_key[i-1](states[...,cursor:cursor+4**i,:])
            keys.append(keys_i)
            cursor += 4**i

        keys = torch.cat(keys, dim=-2)
        return keys

    def to_location_distribution(self, scores, target_depth=-1):
        '''
        scores: batch_size * seq_len * n_nodes / 4 * 4
        scores is sorted according to node_id
        convert the output of forward without target to the log distribution for the depth
        if the depth is -1, the depth is the maximum depth of the tree
        if the depth is 0, return distributions of all layers (i.e., list of batch_size * seq_len * num_nodes at the depth)
        '''
        assert scores.shape[-2] == len(self.tree.get_all_nodes()) - len(self.tree.get_leafs()), "the range of the input distribution should be should be {}, but it is {}".format(len(self.tree.get_all_nodes()) - len(self.tree.get_leafs()), scores.shape[-2])
        if not self.training and self.pre_training:
            print("WARNING: PRE-TRAINING_MODE")
        # scores is sorted according to node.id
        def get_log_distribution_at_depth(scores, depth):

            if self.is_consistent:
                scores = F.log_softmax(scores, dim=-1)
                # when evaluation_mode, conducting consistent sampling for all locations; it generates probability distribution in all layers
                # that is, this generates Pr(location|depth)
                if (not self.training) or self.pre_training:
                # if True:
                    distribution = torch.zeros(*scores.shape[:-2], 4**depth).to(scores.device)
                    for i in range(depth):
                        ids = list(range(sum([4**depth_ for depth_ in range(0,i)]), sum([4**depth_ for depth_ in range(0,i+1)])))
                        score_at_depth_i = scores[...,ids,:].view(*scores.shape[:-2],-1).repeat_interleave(4**(depth-i-1), dim=-1)
                        distribution += score_at_depth_i
                    hidden_ids = list(range(sum([4**depth_ for depth_ in range(0,i+1)])-1, sum([4**depth_ for depth_ in range(0,i+2)])-1))
                    node_ids = [self.tree.hidden_id_to_node_id[id] for id in hidden_ids]
                    distribution = distribution[...,[id-node_ids[0] for id in node_ids]]
                    return distribution

            hidden_ids = list(range(scores.shape[-1] * scores.shape[-2]))
            node_ids = [self.tree.hidden_id_to_node_id[id] for id in hidden_ids]
            scores = scores.view(*scores.shape[:-2], -1)[..., [id-node_ids[0] for id in node_ids]]
            ids = list(range(sum([4**depth_ for depth_ in range(1,depth)]), sum([4**depth_ for depth_ in range(1,depth+1)])))
            if not self.is_consistent:
                # in the case of not consistent, this generates probability distribution on the depth layer
                distribution = F.log_softmax(scores[..., ids], dim=-1)
            else:
                # when training_mode, we consider that each node generates probability distribution on the 4 children nodes, i.e., P(child|parent)
                distribution = scores[..., ids]

            return distribution

        if target_depth == 0:
            return [get_log_distribution_at_depth(scores, depth) for depth in range(1, self.tree.max_depth+1)]
        elif target_depth == -1:
            target_depth = self.tree.max_depth

        return get_log_distribution_at_depth(scores, target_depth)

    def remove_class_to_query(self):
        # self.class_to_query.requires_grad_(False)
        self.pre_training = False
        del self.class_to_query

# class ResidualQuadTreeNetwork(BaseQuadTreeNetwork):
#     def __init__(self, n_locations, memory_dim, hidden_dim, n_classes, activate, multilayer=False, is_consistent=False):
#         super().__init__(n_locations, memory_dim, hidden_dim, n_classes, activate, is_consistent)
#         if multilayer:
#             # self.linears = nn.ModuleList([nn.Sequential(nn.Linear(self.memory_dim, 8*self.memory_dim), self.activate, nn.Linear(8*self.memory_dim, 4*self.memory_dim)) for _ in range(self.tree.max_depth)])
#             self.linears = nn.ModuleList([nn.Sequential(nn.Linear(self.memory_dim, 4*self.memory_dim), self.activate) for _ in range(self.tree.max_depth)])
#         else:
#             # self.linears = nn.ModuleList([nn.Linear(self.memory_dim, 4*self.memory_dim) for _ in range(self.tree.max_depth)])
#             self.linears = nn.ModuleList([nn.Embedding(4, self.memory_dim) for _ in range(self.tree.max_depth)])
#         self.input_dim = hidden_dim
#         # state_to_key is the standard MLP
#         # self.state_to_key = nn.Sequential(nn.Linear(self.memory_dim, self.memory_dim), self.activate, nn.Linear(self.memory_dim, self.memory_dim))
#         # self.state_to_key = nn.Linear(self.memory_dim, self.memory_dim)
#         self.state_to_key = nn.ModuleList([nn.Linear(self.memory_dim, self.memory_dim) for _ in range(self.tree.max_depth)])
#         # self.state_to_key = nn.ModuleList([nn.Sequential(nn.Linear(self.memory_dim, self.memory_dim), nn.Linear(self.memory_dim, self.memory_dim)) for _ in range(self.tree.max_depth)])
#         self.root_value = nn.Embedding(1, self.memory_dim)

#     def make_states(self, shape):
#         '''
#         root_state: batch_size * seq_len * memory_dim
#         output: batch_size * seq_len * (4^depth+4^depth-1+...+4) * memory_dim
#         From the root_state, recursively apply the linear operation to the state until it reaches the maximum depth
#         because the linear layer is (memory_dim, 4*memory_dim), each layer has 4 times parameters (nodes) than the previous layer
#         '''
#         # print(shape)
#         device = self.root_value.weight.device
#         states = []
#         input = torch.tensor([0,1,2,3], device=device).repeat(shape[0], shape[1], 1)
#         vectors = [linear(input) for linear in self.linears]
#         ith_state = self.root_value(torch.zeros(*shape[:-1], device=device).long())
#         root_state = ith_state
#         for i, vector in enumerate(vectors):
#             ith_state = ith_state.repeat_interleave(4, dim=-2).view(*shape[:-1], -1, self.memory_dim) + vector.repeat(*[1]*len(shape[:-1]),4**(i),1).view(*shape[:-1], -1, self.memory_dim)
#             # print(i, ith_state[0][0])
#             states.append(ith_state)
#         states = torch.concat(states, dim=-2)
#         return states

class LinearQuadTreeNetwork(BaseQuadTreeNetwork):
    def __init__(self, n_locations, memory_dim, hidden_dim, n_classes, activate, multilayer=False, is_consistent=False):
        super().__init__(n_locations, memory_dim, hidden_dim, n_classes, activate, is_consistent)
        if multilayer:
            self.linears = nn.ModuleList([nn.Sequential(nn.Linear(self.memory_dim, 8*self.memory_dim), self.activate, nn.Linear(8*self.memory_dim, 4*self.memory_dim)) for _ in range(self.tree.max_depth)])
            # self.linears = nn.ModuleList([nn.Sequential(nn.Linear(self.memory_dim, 4*self.memory_dim), self.activate) for _ in range(self.tree.max_depth)])
        else:
            self.linears = nn.ModuleList([nn.Linear(self.memory_dim, 4*self.memory_dim) for _ in range(self.tree.max_depth)])
        self.input_dim = hidden_dim
        # state_to_key is the standard MLP
        # self.state_to_key = nn.Sequential(nn.Linear(self.memory_dim, self.memory_dim), self.activate, nn.Linear(self.memory_dim, self.memory_dim))
        # self.state_to_key = nn.Linear(self.memory_dim, self.memory_dim)
        self.state_to_key = nn.ModuleList([nn.Linear(self.memory_dim, self.memory_dim) for _ in range(self.tree.max_depth)])
        # self.state_to_key = nn.ModuleList([nn.Sequential(nn.Linear(self.memory_dim, self.memory_dim), nn.Linear(self.memory_dim, self.memory_dim)) for _ in range(self.tree.max_depth)])
        self.root_value = nn.Embedding(1, self.memory_dim)

    def make_states(self, shape):
        '''
        root_state: batch_size * seq_len * memory_dim
        output: batch_size * seq_len * (4^depth+4^depth-1+...+4) * memory_dim
        From the root_state, recursively apply the linear operation to the state until it reaches the maximum depth
        because the linear layer is (memory_dim, 4*memory_dim), each layer has 4 times parameters (nodes) than the previous layer
        '''
        # print(shape)
        device = self.root_value.weight.device
        states = []
        ith_state = self.root_value(torch.zeros(*shape[:-1], device=device).long())
        for i, linear in enumerate(self.linears):
            # residual = root_state.repeat_interleave(4**(i+1), dim=-2).view(*shape[:-1], -1, self.memory_dim)
            # residual = ith_state.repeat(*[1]*len(shape[:-1]), 4, 1).view(*shape[:-1], -1, self.memory_dim)
            ith_state = linear(ith_state).view(ith_state.shape[0], ith_state.shape[1], -1, self.memory_dim)
            states.append(ith_state)
        states = torch.concat(states, dim=-2)

        return states

# in this class, location embedding comes from the tconvs with input of the self.root_value(1)
# this class requires privtree
class FullLinearQuadTreeNetwork(LinearQuadTreeNetwork):
# class FullLinearQuadTreeNetwork(ResidualQuadTreeNetwork):
    def __init__(self, n_locations, memory_dim, hidden_dim, location_embedding_dim, privtree, activate, multilayer, is_consistent):
        # n_classes = len(privtree.get_leafs())
        self.location_embedding_dim = location_embedding_dim
        n_classes = len(privtree.merged_leafs)
        self.n_locations = n_locations
        super().__init__(n_locations, memory_dim, hidden_dim, n_classes, activate, multilayer, is_consistent)
        self.privtree = privtree
        self.root_value = nn.Embedding(4, self.memory_dim)
        # 0 -> states, 1 -> start value, 2 -> ignore value, 3 -> location embedding
        self.leaf_ids = self.privtree.get_leaf_ids_in_tree(self.tree)
        # self.hidden_ids = torch.tensor([self.tree.node_id_to_hidden_id[node_id] for node_id in self.leaf_ids])
        self.node_ids = torch.tensor(self.privtree.get_leaf_ids_in_tree(self.tree))
        # self.class_to_query = nn.Linear(location_embedding_dim, self.memory_dim)
        self.class_to_query = nn.ModuleList([nn.Linear(location_embedding_dim, self.memory_dim) for _ in range(self.tree.max_depth)])
        # self.class_to_query = nn.Sequential(nn.Linear(location_embedding_dim, self.memory_dim), self.activate, nn.Linear(self.memory_dim, self.memory_dim))
        # self.state_to_location_embedding = nn.Linear(self.memory_dim, location_embedding_dim)
        self.state_to_location_embedding = nn.ModuleList([nn.Linear(self.memory_dim, location_embedding_dim) for _ in range(self.tree.max_depth)])
        # self.state_to_location_embedding = nn.ModuleList([nn.Sequential(nn.Linear(self.memory_dim, self.memory_dim), nn.Linear(self.memory_dim, location_embedding_dim)) for _ in range(self.tree.max_depth)])

        self.hidden_to_query_ = nn.ModuleList([nn.Linear(hidden_dim, self.memory_dim) for _ in range(self.tree.max_depth)])
        # self.hidden_to_query_ = nn.ModuleList([nn.Sequential(nn.Linear(hidden_dim, self.memory_dim), nn.Linear(self.memory_dim, self.memory_dim)) for _ in range(self.tree.max_depth)])

        # self.location_linears = nn.ModuleList([nn.Linear(self.memory_dim, 4*self.memory_dim) for _ in range(self.tree.max_depth)])

    def compute_scores(self, querys):
        '''
        scores: batch_size * (seq_len) * (4^depth+4^depth-1+...+4)
        '''
        keys = self.make_keys(querys[0].shape)

        scores = []

        for key, query in zip(keys, querys):
            original_shape = query.shape
            # matmal key and field
            query = query.view(original_shape[0] * original_shape[1], -1, self.memory_dim)
            key = key.view(original_shape[0] * original_shape[1], -1, self.memory_dim)
            # normaize query and key
            # query = F.normalize(query, p=2, dim=-1)
            # key = F.normalize(key, p=2, dim=-1)
            scores.append(torch.bmm(key, query.transpose(-2,-1)).view(*original_shape[:-1], -1, 4)/np.sqrt(self.memory_dim))

        scores = torch.cat(scores, dim=-2)
        return scores

    def hidden_to_query(self, hidden):
       # for pre_training
        if hasattr(self, "class_to_query"):
            # hidden: batch_size * seq_len * n_classes
            node_embeddings = self.location_embedding(self.node_ids.to(hidden.device), True) # n_classes * memory_dim_
            node_embeddings_ = torch.zeros(self.n_classes, self.location_embedding_dim, device=hidden.device)
            for i, leafs in enumerate(self.privtree.merged_leafs):
                for leaf in leafs:
                    index = self.leaf_ids.index(leaf.id)
                    node_embeddings_[i] += node_embeddings[index]

            location_embeddings = hidden.matmul(node_embeddings_) # batch_size * memory_dim

            querys = []
            for linear in self.class_to_query:
                querys.append(linear(location_embeddings))
        else:
            querys = []
            for linear in self.hidden_to_query_:
                querys.append(linear(hidden))
        # output hidden: batch_size * seq_len * memory_dim
        return querys

    def location_embedding(self, location, is_node_id=False):
        # virtual input for grad_sample
        batch_size = location.shape[0]
        start_value = self.root_value(torch.ones(batch_size, device=location.device).long()).view(batch_size, -1, 1, self.memory_dim)
        ignore_value = self.root_value(torch.ones(batch_size, device=location.device).long()*2).view(batch_size, -1, 1, self.memory_dim)
        states = self.make_states_for_location([batch_size, 1, 1])

        location_embeddings = []
        curosor = 0
        for i in range(1,self.tree.max_depth+1):
            location_embeddings.append(self.state_to_location_embedding[i-1](states[...,curosor:curosor+4**i,:]))
            curosor += 4**i
        location_embeddings = torch.cat(location_embeddings, dim=-2)

        start_embeddings = self.state_to_location_embedding[-1](start_value)
        ignore_embeddings = self.state_to_location_embedding[-1](ignore_value)
        # add the start value to states
        location_embeddings = torch.cat([location_embeddings, start_embeddings], dim=-2)
        # add the ignore value to location_embeddings
        location_embeddings = torch.cat([location_embeddings, ignore_embeddings], dim=-2)

        if not is_node_id:
            # input is a state, which is in the order of from upper left to lower right
            # but we need to access to location embedding which is sorted according to node.id
            location_ = []
            for loc in location.view(-1).tolist():
                if loc < self.n_locations:
                    location_.append(self.tree.state_to_node_id_path(loc)[-1]-1)
                else:
                    location_.append(loc)
            location = location_
        else:
            location = location -1
            # because the first node is the root node and make_states does not include the root node
            location = location.view(-1).tolist()
        # states = states[list(range(batch_size)), ..., location, :]
        # location_embeddings = self.state_to_location_embedding(states).view(batch_size, -1)
        # return location_embeddings
        return location_embeddings[list(range(batch_size)), ..., location, :].view(batch_size, -1)

    def make_keys(self, shape):
        '''
        convert state to key by self.state_to_key (MLP)
        shaep should be [batch_size, seq_len, memory_dim]
        output: batch_size * seq_len * (4^depth+4^depth-1+...+4) * memory_dim
        the key is aligned with the order of the node_id in the tree
        '''
        states = self.make_states(shape)

        cursor = 0
        keys = []
        for i in range(1,1+self.tree.max_depth):
            keys_i = self.state_to_key[i-1](states[...,cursor:cursor+4**i,:])
            keys.append(keys_i)
            cursor += 4**i

        return keys

    def make_states_for_location(self, shape):
        '''
        root_state: batch_size * seq_len * memory_dim
        output: batch_size * seq_len * (4^depth+4^depth-1+...+4) * memory_dim
        From the root_state, recursively apply the linear operation to the state until it reaches the maximum depth
        because the linear layer is (memory_dim, 4*memory_dim), each layer has 4 times parameters (nodes) than the previous layer
        '''
        # print(shape)
        device = self.root_value.weight.device
        states = []
        ith_state = self.root_value(3*torch.ones(*shape[:-1], device=device).long())
        # for i, linear in enumerate(self.location_linears):
        for i, linear in enumerate(self.linears):
            # residual = root_state.repeat_interleave(4**(i+1), dim=-2).view(*shape[:-1], -1, self.memory_dim)
            # residual = ith_state.repeat(*[1]*len(shape[:-1]), 4, 1).view(*shape[:-1], -1, self.memory_dim)
            residual = ith_state.repeat_interleave(4, dim=-2).view(*shape[:-1], -1, self.memory_dim)
            ith_state = linear(ith_state).view(ith_state.shape[0], ith_state.shape[1], -1, self.memory_dim) + residual
            states.append(ith_state)
        states = torch.concat(states, dim=-2)

        return states

def guide_to_model(type):
    if type == "hiemrnet":
        return FullLinearQuadTreeNetwork, MetaGRUNet
    elif type == "baseline":
        return MetaNetwork, MetaGRUNet


import logging
import sys

import pytest
import imgui
from imgui import ImGuiError

class _TestException(Exception):
    pass

@pytest.fixture
def context():
    ctx = imgui.get_current_context()
    if ctx is not None:
        imgui.destroy_context(ctx)
    ctx = imgui.create_context()
    io = imgui.get_io()
    io.delta_time = 1.0 / 60.0
    io.display_size = 300, 300

    # setup default font
    io.fonts.get_tex_data_as_rgba32()
    io.fonts.add_font_default()
    io.fonts.texture_id = 0  # set any texture ID to avoid segfaults
    return ctx

@pytest.fixture
def frame(context):
    imgui.new_frame()
    yield
    try:
        imgui.render()
    except ImGuiError:
        try:
            imgui.end_frame()
        except ImGuiError:
            pass

# ------- BEGIN/END ----------
def test_begin_okay(frame):
    imgui.begin("Example: empty window")
    imgui.end()

def test_begin_with(frame):
    with imgui.begin("Example: empty window") as window:
        assert isinstance(window.expanded, bool)
        assert isinstance(window.opened, bool)

def test_begin_with_exception(frame):
    # the real test is that the frame cleanup doesn't crash
    with pytest.raises(_TestException):
        with imgui.begin("Example: empty window"):
            raise _TestException

def test_begin_unpacking(frame):
    expanded, opened = imgui.begin("Example: empty window")
    assert isinstance(expanded, bool)
    assert isinstance(opened, bool)
    imgui.end()

def test_begin_equality(frame):
    window = imgui.begin("Example: empty window")
    assert window == window
    assert window == tuple(window)
    imgui.end()

# ------- BEGIN_CHILD/END_CHILD ----------
def test_child_okay(frame):
    imgui.begin("Example: child region")

    imgui.begin_child("region", 150, -50, border=True)
    imgui.text("inside region")
    imgui.end_child()

    imgui.text("outside region")
    imgui.end()

def test_child_with(frame):
    with imgui.begin("Example: child region"):
        with imgui.begin_child("region", 150, -50, border=True):
            imgui.text("inside region")
        imgui.text("outside region")

def test_child_with_exception(frame):
    # the real test is that the frame cleanup doesn't crash
    with pytest.raises(_TestException):
        with imgui.begin("Example: child region"):
            with imgui.begin_child("region", 150, -50, border=True):
                raise _TestException
            imgui.text("outside region")

def test_child_as_bool(frame):
    imgui.begin("Example: child region")
    child = imgui.begin_child("region", 150, -50, border=True)
    assert bool(child) is child.visible
    assert child == child
    assert child == bool(child)
    imgui.end_child()
    imgui.end()

# ------- BEGIN_TOOLTIP/END_TOOLTIP ----------
def test_tooltip_okay(frame):
    imgui.begin("Example: tooltip")
    imgui.button("Click me!")
    if imgui.is_item_hovered():
        imgui.begin_tooltip()
        imgui.text("This button is clickable.")
        imgui.end_tooltip()
    imgui.end()

def test_tooltip_with(frame):
    with imgui.begin("Example: tooltip"):
        imgui.button("Click me!")
        if imgui.is_item_hovered():
            with imgui.begin_tooltip():
                imgui.text("This button is clickable.")

def test_tooltip_with_exception(frame):
    # the real test is that the frame cleanup doesn't crash
    with pytest.raises(_TestException):
        with imgui.begin("Example: tooltip"):
            imgui.button("Click me!")
            with imgui.begin_tooltip():
                raise _TestException

# ------- BEGIN_MAIN_MENU_BAR/END_MAIN_MENU_BAR ----------
def test_main_menu_bar_okay(frame):
    if imgui.begin_main_menu_bar():
        imgui.end_main_menu_bar()
    else:
        assert False

def test_main_menu_bar_with(frame):
    with imgui.begin_main_menu_bar() as menu:
        assert menu.opened is True

def test_main_menu_bar_with_exception(frame):
    # the real test is that the frame cleanup doesn't crash
    with pytest.raises(_TestException):
        with imgui.begin_main_menu_bar():
            raise _TestException

def test_main_menu_bar_as_bool(frame):
    menu = imgui.begin_main_menu_bar()
    assert isinstance(menu.opened, bool)
    assert bool(menu) is menu.opened is True
    assert menu == menu
    assert menu == bool(menu)
    imgui.end_main_menu_bar()

# ------- BEGIN_MENU_BAR/END_MENU_BAR ----------
def test_menu_bar_okay(frame):
    flags = imgui.WINDOW_MENU_BAR
    imgui.begin("Child Window - File Browser", flags=flags)

    if imgui.begin_menu_bar():
        imgui.end_menu_bar()
    else:
        assert False
    imgui.end()

def test_menu_bar_with(frame):
    flags = imgui.WINDOW_MENU_BAR
    with imgui.begin("Child Window - File Browser", flags=flags):
        with imgui.begin_menu_bar():
            pass

def test_menu_bar_with_exception(frame):
    # the real test is that the frame cleanup doesn't crash
    with pytest.raises(_TestException):
        flags = imgui.WINDOW_MENU_BAR
        with imgui.begin("Child Window - File Browser", flags=flags):
            with imgui.begin_menu_bar():
                raise _TestException

def test_menu_bar_as_bool(frame):
    flags = imgui.WINDOW_MENU_BAR
    imgui.begin("Child Window - File Browser", flags=flags)
    menu = imgui.begin_menu_bar()
    assert isinstance(menu.opened, bool)
    assert bool(menu) is menu.opened is True
    assert menu == menu
    assert menu == bool(menu)
    imgui.end_menu_bar()
    imgui.end()

# ------- BEGIN_TAB_BAR/END_TAB_BAR ----------
def test_tab_bar_okay(frame):
    imgui.begin("Example Tab Bar")
    if imgui.begin_tab_bar("MyTabBar"):
        imgui.end_tab_bar()
    imgui.end()

def test_tab_bar_with(frame):
    with imgui.begin("Example Tab Bar"):
        with imgui.begin_tab_bar("MyTabBar"):
            pass

def test_tab_bar_with_exception(frame):
    # the real test is that the frame cleanup doesn't crash
    with pytest.raises(_TestException):
        with imgui.begin("Example Tab Bar"):
            with imgui.begin_tab_bar("MyTabBar"):
                raise _TestException

def test_tab_bar_as_bool(frame):
    imgui.begin("Example Tab Bar")
    tab_bar = imgui.begin_tab_bar("MyTabBar")
    assert bool(tab_bar) is tab_bar.opened is True
    assert tab_bar == tab_bar
    assert tab_bar == bool(tab_bar)
    imgui.end_tab_bar()
    imgui.end()

# ------- BEGIN_TAB_ITEM/END_TAB_ITEM ----------
def test_tab_item_okay(frame):
    imgui.begin("Example Tab Bar")
    if imgui.begin_tab_bar("MyTabBar"):
        if imgui.begin_tab_item("Item 1")[0]:
            imgui.text("Here is the tab content!")
            imgui.end_tab_item()
        imgui.end_tab_bar()
    imgui.end()

def test_tab_item_with(frame):
    with imgui.begin("Example Tab Bar"):
        with imgui.begin_tab_bar("MyTabBar") as tab_bar:
            if tab_bar.opened:
                with imgui.begin_tab_item("Item 1") as item:
                    if item.selected:
                        pass

def test_tab_item_with_exception(frame):
    # the real test is that the frame cleanup doesn't crash
    with pytest.raises(_TestException):
        with imgui.begin("Example Tab Bar"):
            with imgui.begin_tab_bar("MyTabBar") as tab_bar:
                if tab_bar.opened:
                    with imgui.begin_tab_item("Item 1") as item:
                        raise _TestException

def test_tab_item_as_bool(frame):
    imgui.begin("Example Tab Bar")
    tab_bar = imgui.begin_tab_bar("MyTabBar")
    assert tab_bar.opened
    item = imgui.begin_tab_item("Item 1")
    assert bool(item) is item.selected is True
    assert item.opened is False
    assert item == item
    assert item == tuple(item)
    imgui.end_tab_item()
    imgui.end_tab_bar()
    imgui.end()

# ------- BEGIN_MENU/END_MENU ----------
def test_menu_okay(frame):
    flags = imgui.WINDOW_MENU_BAR
    imgui.begin("Child Window - File Browser", flags=flags)

    if imgui.begin_menu_bar():
        if imgui.begin_menu('File'):
            imgui.menu_item('Close')
            imgui.end_menu()
        imgui.end_menu_bar()
    imgui.end()

def test_menu_with(frame):
    flags = imgui.WINDOW_MENU_BAR
    with imgui.begin("Child Window - File Browser", flags=flags):
        with imgui.begin_menu_bar() as menu_bar:
            if menu_bar.opened:
                with imgui.begin_menu('File'):
                    pass

def test_menu_with_exception(frame):
    # the real test is that the frame cleanup doesn't crash
    with pytest.raises(_TestException):
        flags = imgui.WINDOW_MENU_BAR
        with imgui.begin("Child Window - File Browser", flags=flags):
            with imgui.begin_menu_bar() as menu_bar:
                if menu_bar.opened:
                    with imgui.begin_menu('File'):
                        raise _TestException

def test_menu_as_bool(frame):
    flags = imgui.WINDOW_MENU_BAR
    imgui.begin("Child Window - File Browser", flags=flags)
    menu_bar = imgui.begin_menu_bar()
    assert menu_bar.opened
    menu = imgui.begin_menu('File')
    assert isinstance(menu.opened, bool)
    assert bool(menu) is menu.opened is False
    assert menu == menu
    assert menu == bool(menu)
    imgui.end_menu_bar()
    imgui.end()

# ------- BEGIN_POPUP{,_CONTEXT_ITEM,_CONTEXT_WINDOW,_CONTEXT_VOID}/END_POPUP ----------
def test_popup_okay(frame):
    imgui.begin("Example: simple popup")
    imgui.open_popup("select-popup")
    if imgui.begin_popup("select-popup"):
        imgui.end_popup()
    imgui.end()

def test_popup_with(frame):
    with imgui.begin("Example: simple popup"):
        imgui.open_popup("select-popup")
        with imgui.begin_popup("select-popup"):
            pass

def test_popup_with_exception(frame):
    # the real test is that the frame cleanup doesn't crash
    with pytest.raises(_TestException):
        with imgui.begin("Example: simple popup"):
            imgui.open_popup("select-popup")
            with imgui.begin_popup("select-popup"):
                raise _TestException

def test_popup_as_bool(frame):
    imgui.begin("Example: simple popup")
    imgui.open_popup("select-popup")
    popup = imgui.begin_popup("select-popup")
    assert isinstance(popup.opened, bool)
    assert bool(popup) is popup.opened is True
    assert popup == popup
    assert popup == bool(popup)
    imgui.end_popup()
    imgui.end()

def test_popup_context_item_isinstance(frame):
    imgui.begin("Example: popup context view")
    imgui.text("Right-click to set value.")
    item = imgui.begin_popup_context_item("Item Context Menu")
    assert isinstance(item, imgui.core._BeginEndPopup)
    assert item.opened is False
    assert item == item
    assert item == bool(item)
    imgui.end()

def test_popup_context_window_isinstance(frame):
    imgui.begin("Example: popup context window")
    window = imgui.begin_popup_context_window()
    assert isinstance(window, imgui.core._BeginEndPopup)
    assert window.opened is False
    assert window == window
    assert window == bool(window)
    imgui.end()

def test_popup_context_void_isinstance(frame):
    window = imgui.begin_popup_context_void()
    assert isinstance(window, imgui.core._BeginEndPopup)
    assert window.opened is False
    assert window == window
    assert window == bool(window)

# ------- BEGIN_POPUP_MODAL/END_POPUP_MODAL ----------
def test_popup_modal_okay(frame):
    imgui.begin("Example: simple popup modal")
    imgui.open_popup("select-popup")
    if imgui.begin_popup_modal("select-popup")[0]:
        imgui.end_popup()
    imgui.end()

def test_popup_modal_with(frame):
    with imgui.begin("Example: simple popup modal"):
        imgui.open_popup("select-popup")
        with imgui.begin_popup_modal("select-popup"):
            pass

def test_popup_modal_with_exception(frame):
    # the real test is that the frame cleanup doesn't crash
    with pytest.raises(_TestException):
        with imgui.begin("Example: simple popup modal"):
            imgui.open_popup("select-popup")
            with imgui.begin_popup_modal("select-popup"):
                raise _TestException

def test_popup_modal_as_bool(frame):
    imgui.begin("Example: simple popup modal")
    imgui.open_popup("select-popup")
    popup = imgui.begin_popup_modal("select-popup")
    assert popup.opened is True
    assert popup.visible is False
    opened, visible = popup
    assert opened is popup.opened is popup[0]
    assert visible is popup.visible is popup[1]
    assert popup == popup
    assert popup == tuple(popup)
    imgui.end_popup()
    imgui.end()

# ------- BEGIN_DRAG_DROP_SOURCE/END_DRAG_DROP_SOURCE ----------
def test_drag_drop_source_okay(frame):
    imgui.begin("Example: drag and drop")
    imgui.button('source')
    if imgui.begin_drag_drop_source():
        imgui.set_drag_drop_payload('itemtype', b'payload')
        imgui.end_drag_drop_source()
    imgui.end()

def test_drag_drop_source_with(frame):
    with imgui.begin("Example: drag and drop"):
        imgui.button('source')
        with imgui.begin_drag_drop_source():
            pass

def test_drag_drop_source_with_exception(frame):
    # the real test is that the frame cleanup doesn't crash
    with pytest.raises(_TestException):
        with imgui.begin("Example: drag and drop"):
            imgui.button('source')
            with imgui.begin_drag_drop_source():
                raise _TestException

def test_drag_drop_source_as_bool(frame):
    imgui.begin("Example: drag and drop")
    imgui.button('source')
    src = imgui.begin_drag_drop_source()
    assert bool(src) is src.dragging is False
    assert src == src
    assert src == bool(src)
    imgui.end()

# ------- BEGIN_DRAG_DROP_TARGET/END_DRAG_DROP_TARGET ----------
def test_drag_drop_target_okay(frame):
    imgui.begin("Example: drag and drop")
    imgui.button('dest')
    if imgui.begin_drag_drop_target():
        payload = imgui.accept_drag_drop_payload('itemtype')
        imgui.end_drag_drop_target()
    imgui.end()

def test_drag_drop_target_with(frame):
    with imgui.begin("Example: drag and drop"):
        imgui.button('dest')
        with imgui.begin_drag_drop_target():
            pass

def test_drag_drop_target_with_exception(frame):
    # the real test is that the frame cleanup doesn't crash
    with pytest.raises(_TestException):
        with imgui.begin("Example: drag and drop"):
            imgui.button('dest')
            with imgui.begin_drag_drop_target():
                raise _TestException

def test_drag_drop_target_as_bool(frame):
    imgui.begin("Example: drag and drop")
    imgui.button('dest')
    target = imgui.begin_drag_drop_target()
    assert bool(target) is target.hovered is False
    assert target == target
    assert target == bool(target)
    imgui.end()

# ------- BEGIN_GROUP/END_GROUP ----------
def test_group_okay(frame):
    imgui.begin("Example: item groups")
    imgui.begin_group()
    imgui.text("First group (buttons):")
    imgui.end_group()
    imgui.end()

def test_group_with(frame):
    with imgui.begin("Example: item groups"):
        with imgui.begin_group():
            imgui.text("First group (buttons):")

def test_group_with_exception(frame):
    # the real test is that the frame cleanup doesn't crash
    with pytest.raises(_TestException):
        with imgui.begin("Example: item groups"):
            with imgui.begin_group():
                raise _TestException

# ------- BEGIN_LIST_BOX/END_LIST_BOX ----------
def test_list_box_okay(frame):
    imgui.begin("Example: custom listbox")
    if imgui.begin_list_box("List", 200, 100):
        imgui.selectable("Selected", True)
        imgui.selectable("Not Selected", False)
        imgui.end_list_box()
    imgui.end()

def test_list_box_with(frame):
    with imgui.begin("Example: custom listbox"):
        with imgui.begin_list_box("List", 200, 100):
            pass

def test_list_box_with_exception(frame):
    # the real test is that the frame cleanup doesn't crash
    with pytest.raises(_TestException):
        with imgui.begin("Example: custom listbox"):
            with imgui.begin_list_box("List", 200, 100):
                raise _TestException

def test_list_box_as_bool(frame):
    imgui.begin("Example: custom listbox")
    list_box = imgui.begin_list_box("List", 200, 100)
    assert bool(list_box) is list_box.opened is True
    assert list_box == list_box
    assert list_box == bool(list_box)
    imgui.end_list_box()
    imgui.end()

# ------- BEGIN_TABLE/END_TABLE ----------
def test_table_okay(frame):
    imgui.begin("Example: table")
    if imgui.begin_table("data", 2):
        imgui.end_table()
    imgui.end()

def test_table_with(frame):
    with imgui.begin("Example: table"):
        with imgui.begin_table("data", 2):
            pass

def test_table_with_exception(frame):
    # the real test is that the frame cleanup doesn't crash
    with pytest.raises(_TestException):
        with imgui.begin("Example: table"):
            with imgui.begin_table("data", 2):
                raise _TestException

def test_table_as_bool(frame):
    imgui.begin("Example: table")
    table = imgui.begin_table("data", 2)
    assert bool(table) is table.opened is True
    assert table == table
    assert table == bool(table)
    imgui.end_table()
    imgui.end()


# https://leetcode.com/problems/product-of-array-except-self/
# 3/28/2023

# 238. Product of Array Except Self
# Given an integer array nums, return an array answer such that answer[i] is equal
# to the product of all the elements of nums except nums[i].
# The product of any prefix or suffix of nums is guaranteed to fit in a 32-bit integer.
# You must write an algorithm that runs in O(n) time and without using the division operation.
from typing import List

class Solution:
    def productExceptSelf(self, nums: List[int]) -> List[int]:
        prefix_product = nums.copy()
        suffix_product = nums.copy()
        for i in range(1,len(nums)):
            prefix_product[i] = prefix_product[i-1] * nums[i]
        for i in range(len(nums)-2,-1,-1):
            suffix_product[i] = suffix_product[i+1] * nums[i]
        answer = []
        answer.append(suffix_product[1])
        for i in range(1,len(nums)-1):
            answer.append(prefix_product[i-1]*suffix_product[i+1])
        answer.append(prefix_product[len(nums)-2])
        return answer


import itertools

import os
from collections import Counter

def process_frames(p,q,t):
    average_ensemble_sum=[]
    for i in range(0,48):
        average_ensemble_sum.append(0)
    ensemble_set=[]
    # Creating 48 ensembles of the frame set,ensemble_set containing 48 ensembles
    ensemble_set=sampling(p,q)

    # Creating 100 video volumes in each ensemble,vol_set containg 48 lists,
    # each list contains 100 video volumes
    vol_set=[]
    gray_sum=[]
    gray_sum_l=[]
    gray_sum_s=[]
    for i in range(0,48):
        vol_set.append(video_volume(ensemble_set[i],q))

    for i in range(0,48):
            for j in range(0,100):
                for k in range(0,q):
                    vol_set[i][j][k]=list(itertools.chain.from_iterable(vol_set[i][j][k]))
                    gray_sum.append(sum(vol_set[i][j][k]))
    s=0
    e=q
    for i in range(0,4800):
        gray_sum_l.append(max(gray_sum[s:e]))
        gray_sum_s.append(min(gray_sum[s:e]))
        s+=q
        e+=q
    re_list=[]
    re_list.append(gray_sum_l)
    re_list.append(gray_sum_s)
    return re_list

#Sampling function this will create 48 samples of incoming video sequence 640*480
def sampling(p,q):
    total_frames=q
    ensemble_row=6
    ensemble_col=8
    frame_ensemble=[]
    ensemble_set=[]

    args1=[0,80,total_frames,p]
    args2=[80,160,total_frames,p]
    args3=[160,240,total_frames,p]
    args4=[240,320,total_frames,p]
    args5=[320,400,total_frames,p]
    args6=[400,480,total_frames,p]
    args_list=[args1,args2,args3,args4,args5,args6]

    # parallising part
    # ensemble_set = f.map(core_ensembling,args_list)
    # ensemble_set=list(itertools.chain.from_iterable(ensemble_set))
    # parallising part

    set1=core_ensembling(args1)
    set2=core_ensembling(args2)
    set3=core_ensembling(args3)
    set4=core_ensembling(args4)
    set5=core_ensembling(args5)
    set6=core_ensembling(args6)

    ensemble_set = list(itertools.chain(set1,set2,set3,set4,set5,set6))

    return ensemble_set

# video_volume function will create 100 video voulme of a ensemble.An ensemble is divided in to 100 video volumes
# 80*80 ensemble is divided into 100 8*8 video volumes
def video_volume(ensemble,q):
        volume_set=[]
        total_frames=q
        volume_row=10
        volume_col=10

        args1=[0,8,0,8,volume_row,volume_col,ensemble,q]
        args2=[8,16,0,8,volume_row,volume_col,ensemble,q]
        args3=[16,24,0,8,volume_row,volume_col,ensemble,q]
        args4=[24,32,0,8,volume_row,volume_col,ensemble,q]
        args5=[32,40,0,8,volume_row,volume_col,ensemble,q]
        args6=[40,48,0,8,volume_row,volume_col,ensemble,q]
        args7=[48,56,0,8,volume_row,volume_col,ensemble,q]
        args8=[56,64,0,8,volume_row,volume_col,ensemble,q]
        args9=[64,72,0,8,volume_row,volume_col,ensemble,q]
        args10=[72,80,0,8,volume_row,volume_col,ensemble,q]
        # args_list=[args1,args2,args3,args4,args5,args6,args7,args8,args9,args10]
        #
        # volume_set = g.map(core_volume,args_list)
        # volume_set=list(itertools.chain.from_iterable(volume_set))

        set1=core_volume(args1)
        set2=core_volume(args2)
        set3=core_volume(args3)
        set4=core_volume(args4)
        set5=core_volume(args5)
        set6=core_volume(args6)
        set7=core_volume(args7)
        set8=core_volume(args8)
        set9=core_volume(args9)
        set10=core_volume(args10)

        volume_set = list(itertools.chain(set1,set2,set3,set4,set5,set6,set7,set8,set9,set10))

        return volume_set

#used in sampling function to do repetative task
def core_ensembling(args):
    ensemble_set=[]
    row_no=args[0]
    max_row=args[1]
    total_frames=args[2]
    p=args[3]
    c1=0
    c2=80
    for j in range(0,8):
        ensemble=[]
        for k in range(0,total_frames):
            frame_ensemble=[]
            for l in range(row_no,max_row):
                frame_ensemble.append(p[k][l][c1:c2])
            ensemble.append(frame_ensemble)
        c1=c1+80
        c2=c2+80
        ensemble_set.append(ensemble)
    return ensemble_set

#used in video_volume function to do repetative task
def core_volume(args):
    v_set=[]
    row_no=args[0]
    max_row=args[1]
    c1=args[2]
    c2=args[3]
    volume_row=args[4]
    volume_col=args[5]
    ensemble=args[6]
    total_frames=args[7]
    for j in range(0,volume_col):
        single_volume=[]
        for k in range(0,total_frames):
            frame_volume=[]
            for l in range(row_no,max_row):
                frame_volume.append(ensemble[k][l][c1:c2])
            single_volume.append(frame_volume)
        v_set.append(single_volume)
        c1=c1+8
        c2=c2+8
    return v_set

def clear_folder():
    dirPath = 'C:\\FlaskVideo\\static\\abnormal_frames'
    fileList = os.listdir(dirPath)
    for fileName in fileList:
     os.remove(dirPath+"/"+fileName)

def mod_function(alist):
    mod_list=[[x] for x in range(0,4800)]
    temp_list=[]
    f_c=0
    for i in range(0,4800):
        count = Counter(alist[i])
        temp_list=count.most_common()
        l=len(temp_list)
        if l<2:
            mod_list[i]=temp_list[0][0]
        if l>1:
            for p in range(0,l):
                mod_list[i].append(temp_list[p][0])
    return mod_list

def stat_range(alist):
    ulist=sorted(alist)
    p=sum(alist)/len(alist)
    return p
def round(alist):
    rounding_factor=1000
    blist=[]
    temp_var=0
    l=len(alist)
    for i in range(0,l):
        rem=alist[i]%rounding_factor
        if rem==0:
            blist.append(alist[i])
        if rem!=0:
            alist[i]=alist[i]-rem
            blist.append(alist[i])
    return blist
def most_Common(lst):
    data = Counter(lst)
    return data.most_common(1)[0][0]


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import argparse
import json
import math
import multiprocessing
import os
import re
import sys
import six
import tensorflow as tf

from tensorflow.contrib.framework.python.ops import variables as contrib_variables
from tensorflow.contrib.learn.python.learn import export_strategy
from tensorflow.contrib.learn.python.learn import learn_runner
from tensorflow.contrib.learn.python.learn.estimators import model_fn as model_fn_lib
from tensorflow.contrib.learn.python.learn.utils import saved_model_export_utils
from tensorflow.python.client import session as tf_session
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.lib.io import file_io
from tensorflow.python.ops import resources
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import variables
from tensorflow.python.saved_model import builder as saved_model_builder
from tensorflow.python.saved_model import signature_def_utils
from tensorflow.python.saved_model import tag_constants
from tensorflow.python.training import saver
from tensorflow.python.util import compat

from . import feature_transforms
from . import feature_analysis

# Constants for the Prediction Graph fetch tensors.
PG_TARGET = 'target'  # from input

PG_REGRESSION_PREDICTED_TARGET = 'predicted'

PG_CLASSIFICATION_FIRST_LABEL = 'predicted'
PG_CLASSIFICATION_FIRST_SCORE = 'probability'
PG_CLASSIFICATION_LABEL_TEMPLATE = 'predicted_%s'
PG_CLASSIFICATION_SCORE_TEMPLATE = 'probability_%s'

class DatalabParser():
  """An arg parser that also prints package specific args with --datalab-help.

  When using Datalab magic's to run this trainer, it prints it's own help menu
  that describes the required options that are common to all trainers. In order
  to print just the options that are unique to this trainer, datalab calls this
  file with --datalab-help.

  This class implements --datalab-help by building a list of help string that only
  includes the unique parameters.
  """

  def __init__(self, epilog=None, datalab_epilog=None):
    self.full_parser = argparse.ArgumentParser(epilog=epilog)
    self.datalab_help = []
    self.datalab_epilog = datalab_epilog

    # Datalab help string
    self.full_parser.add_argument(
        '--datalab-help', action=self.make_datalab_help_action(),
        help='Show a smaller help message for DataLab only and exit')

    # The arguments added here are required to exist by Datalab's "%%ml train" magics.
    self.full_parser.add_argument(
        '--train', type=str, required=True, action='append', metavar='FILE')
    self.full_parser.add_argument(
        '--eval', type=str, required=True, action='append', metavar='FILE')
    self.full_parser.add_argument('--job-dir', type=str, required=True)
    self.full_parser.add_argument(
        '--analysis', type=str,
        metavar='ANALYSIS_OUTPUT_DIR',
        help=('Output folder of analysis. Should contain the schema, stats, and '
              'vocab files. Path must be on GCS if running cloud training. ' +
              'If absent, --schema and --features must be provided and ' +
              'the master trainer will do analysis locally.'))
    self.full_parser.add_argument(
        '--transform', action='store_true', default=False,
        help='If used, input data is raw csv that needs transformation. If analysis ' +
             'is required to run in trainerm this is automatically set to true.')
    self.full_parser.add_argument(
        '--schema', type=str,
        help='Schema of the training csv file. Only needed if analysis is required.')
    self.full_parser.add_argument(
        '--features', type=str,
        help='Feature transform config. Only needed if analysis is required.')

  def make_datalab_help_action(self):
    """Custom action for --datalab-help.

    The action output the package specific parameters and will be part of "%%ml train"
    help string.
    """
    datalab_help = self.datalab_help
    epilog = self.datalab_epilog

    class _CustomAction(argparse.Action):

      def __init__(self, option_strings, dest, help=None):
        super(_CustomAction, self).__init__(
            option_strings=option_strings, dest=dest, nargs=0, help=help)

      def __call__(self, parser, args, values, option_string=None):
        print('\n\n'.join(datalab_help))
        if epilog:
          print(epilog)

        # We have printed all help string datalab needs. If we don't quit, it will complain about
        # missing required arguments later.
        quit()
    return _CustomAction

  def add_argument(self, name, **kwargs):
    # Any argument added here is not required by Datalab, and so is unique
    # to this trainer. Add each argument to the main parser and the datalab helper string.
    self.full_parser.add_argument(name, **kwargs)
    name = name.replace('--', '')
    # leading spaces are needed for datalab's help formatting.
    msg = '  ' + name + ': '
    if 'help' in kwargs:
      msg += kwargs['help'] + ' '
    if kwargs.get('required', False):
      msg += 'Required. '
    else:
      msg += 'Optional. '
    if 'choices' in kwargs:
      msg += 'One of ' + str(kwargs['choices']) + '. '
    if 'default' in kwargs:
      msg += 'default: ' + str(kwargs['default']) + '.'
    self.datalab_help.append(msg)

  def parse_known_args(self, args=None):
    return self.full_parser.parse_known_args(args=args)

def parse_arguments(argv):
  """Parse the command line arguments."""
  parser = DatalabParser(
      epilog=('Note that if using a DNN model, --hidden-layer-size1=NUM, '
              '--hidden-layer-size2=NUM, ..., is also required. '),
      datalab_epilog=("""
  Note that if using a DNN model,
  hidden-layer-size1: NUM
  hidden-layer-size2: NUM
  ...
  is also required. """))

  # HP parameters
  parser.add_argument(
      '--epsilon', type=float, default=0.0005, metavar='R',
      help='tf.train.AdamOptimizer epsilon. Only used in dnn models.')
  parser.add_argument(
      '--l1-regularization', type=float, default=0.0, metavar='R',
      help='L1 term for linear models.')
  parser.add_argument(
      '--l2-regularization', type=float, default=0.0, metavar='R',
      help='L2 term for linear models.')

  # Model parameters
  parser.add_argument(
    '--model', required=True,
    choices=['linear_classification', 'linear_regression', 'dnn_classification', 'dnn_regression'])
  parser.add_argument(
      '--top-n', type=int, default=0, metavar='N',
      help=('For classification problems, the output graph will contain the '
            'labels and scores for the top n classes, and results will be in the form of '
            '"predicted, predicted_2, ..., probability, probability_2, ...". '
            'If --top-n=0, then all labels and scores are returned in the form of '
            '"predicted, class_name1, class_name2,...".'))

  # HP parameters
  parser.add_argument(
      '--learning-rate', type=float, default=0.01, metavar='R',
      help='optimizer learning rate.')

  # Training input parameters
  parser.add_argument(
      '--max-steps', type=int, metavar='N',
      help='Maximum number of training steps to perform. If unspecified, will '
           'honor "max-epochs".')
  parser.add_argument(
      '--max-epochs', type=int, default=1000, metavar='N',
      help='Maximum number of training data epochs on which to train. If '
           'both "max-steps" and "max-epochs" are specified, the training '
           'job will run for "max-steps" or "num-epochs", whichever occurs '
           'first. If early stopping is enabled, training may also stop '
           'earlier.')
  parser.add_argument(
      '--train-batch-size', type=int, default=64, metavar='N',
      help='How many training examples are used per step. If num-epochs is '
           'used, the last batch may not be full.')
  parser.add_argument(
      '--eval-batch-size', type=int, default=64, metavar='N',
      help='Batch size during evaluation. Larger values increase performance '
           'but also increase peak memory usgae on the master node. One pass '
           'over the full eval set is performed per evaluation run.')
  parser.add_argument(
      '--min-eval-frequency', type=int, default=1000, metavar='N',
      help='Minimum number of training steps between evaluations. Evaluation '
           'does not occur if no new checkpoint is available, hence, this is '
           'the minimum. If 0, the evaluation will only happen after training. ')
  parser.add_argument(
      '--early-stopping-num_evals', type=int, default=3,
      help='Automatic training stop after results of specified number of evals '
           'in a row show the model performance does not improve. Set to 0 to '
           'disable early stopping.')
  parser.add_argument(
      '--logging-level', choices=['error', 'warning', 'info'],
      help='The TF logging level. If absent, use info for cloud training '
           'and warning for local training.')

  args, remaining_args = parser.parse_known_args(args=argv[1:])

  # All HP parambeters must be unique, so we need to support an unknown number
  # of --hidden-layer-size1=10 --lhidden-layer-size2=10 ...
  # Look at remaining_args for hidden-layer-size\d+ to get the layer info.

  # Get number of layers
  pattern = re.compile('hidden-layer-size(\d+)')
  num_layers = 0
  for other_arg in remaining_args:
    match = re.search(pattern, other_arg)
    if match:
      if int(match.group(1)) <= 0:
        raise ValueError('layer size must be a positive integer. Was given %s' % other_arg)
      num_layers = max(num_layers, int(match.group(1)))

  # Build a new parser so we catch unknown args and missing layer_sizes.
  parser = argparse.ArgumentParser()
  for i in range(num_layers):
    parser.add_argument('--hidden-layer-size%s' % str(i + 1), type=int, required=True)

  layer_args = vars(parser.parse_args(args=remaining_args))
  hidden_layer_sizes = []
  for i in range(num_layers):
    key = 'hidden_layer_size%s' % str(i + 1)
    hidden_layer_sizes.append(layer_args[key])

  assert len(hidden_layer_sizes) == num_layers
  args.hidden_layer_sizes = hidden_layer_sizes

  return args

def is_linear_model(model_type):
  return model_type.startswith('linear_')

def is_dnn_model(model_type):
  return model_type.startswith('dnn_')

def is_regression_model(model_type):
  return model_type.endswith('_regression')

def is_classification_model(model_type):
  return model_type.endswith('_classification')

def build_feature_columns(features, stats, model_type):
  feature_columns = []
  is_dnn = is_dnn_model(model_type)

  # Supported transforms:
  # for DNN
  #   numerical number
  #   one hot: sparse int column -> one_hot_column
  #   ebmedding: sparse int column -> embedding_column
  #   text: sparse int weighted column -> embedding_column
  # for linear
  #   numerical number
  #   one hot: sparse int column
  #   ebmedding: sparse int column -> hash int
  #   text: sparse int weighted column
  # It is unfortunate that tf.layers has different feature transforms if the
  # model is linear or DNN. This pacakge should not expose to the user that
  # we are using tf.layers.
  for name, transform in six.iteritems(features):
    transform_name = transform['transform']
    source_column = transform['source_column']

    if transform_name in feature_transforms.NUMERIC_TRANSFORMS:
      new_feature = tf.contrib.layers.real_valued_column(name, dimension=1)
    elif (transform_name == feature_transforms.ONE_HOT_TRANSFORM or
          transform_name == feature_transforms.MULTI_HOT_TRANSFORM):
      sparse = tf.contrib.layers.sparse_column_with_integerized_feature(
          name,
          bucket_size=stats['column_stats'][source_column]['vocab_size'])
      if is_dnn:
        new_feature = tf.contrib.layers.one_hot_column(sparse)
      else:
        new_feature = sparse
    elif transform_name == feature_transforms.EMBEDDING_TRANSFROM:
      if is_dnn:
        sparse = tf.contrib.layers.sparse_column_with_integerized_feature(
            name,
            bucket_size=stats['column_stats'][source_column]['vocab_size'])
        new_feature = tf.contrib.layers.embedding_column(
            sparse,
            dimension=transform['embedding_dim'])
      else:
        new_feature = tf.contrib.layers.sparse_column_with_hash_bucket(
            name,
            hash_bucket_size=transform['embedding_dim'],
            dtype=dtypes.int64)
    elif transform_name in feature_transforms.TEXT_TRANSFORMS:
      sparse_ids = tf.contrib.layers.sparse_column_with_integerized_feature(
          name + '_ids',
          bucket_size=stats['column_stats'][source_column]['vocab_size'],
          combiner='sum')
      sparse_weights = tf.contrib.layers.weighted_sparse_column(
          sparse_id_column=sparse_ids,
          weight_column_name=name + '_weights',
          dtype=dtypes.float32)
      if is_dnn:
        new_feature = tf.contrib.layers.one_hot_column(sparse_ids)
        dimension = int(math.log(stats['column_stats'][source_column]['vocab_size'])) + 1
        new_feature = tf.contrib.layers.embedding_column(
            sparse_weights,
            dimension=dimension,
            combiner='sqrtn')
      else:
        new_feature = sparse_weights
    elif (transform_name == feature_transforms.TARGET_TRANSFORM or
          transform_name == feature_transforms.KEY_TRANSFORM):
      continue
    elif transform_name == feature_transforms.IMAGE_TRANSFORM:
      new_feature = tf.contrib.layers.real_valued_column(
          name,
          dimension=feature_transforms.IMAGE_HIDDEN_TENSOR_SIZE)
    else:
      raise ValueError('Unknown transfrom %s' % transform_name)

    feature_columns.append(new_feature)

  return feature_columns

def recursive_copy(src_dir, dest_dir):
  """Copy the contents of src_dir into the folder dest_dir.
  Args:
    src_dir: gsc or local path.
    dest_dir: gcs or local path.
  """

  file_io.recursive_create_dir(dest_dir)
  for file_name in file_io.list_directory(src_dir):
    old_path = os.path.join(src_dir, file_name)
    new_path = os.path.join(dest_dir, file_name)

    if file_io.is_directory(old_path):
      recursive_copy(old_path, new_path)
    else:
      file_io.copy(old_path, new_path, overwrite=True)

def make_prediction_output_tensors(args, features, input_ops, model_fn_ops,
                                   keep_target):
  """Makes the final prediction output layer."""
  target_name = feature_transforms.get_target_name(features)
  key_names = get_key_names(features)

  outputs = {}
  outputs.update({key_name: tf.squeeze(input_ops.features[key_name])
                  for key_name in key_names})

  if is_classification_model(args.model):

    # build maps from ints to the origional categorical strings.
    class_names = read_vocab(args, target_name)
    table = tf.contrib.lookup.index_to_string_table_from_tensor(
        mapping=class_names,
        default_value='UNKNOWN')

    # Get the label of the input target.
    if keep_target:
      input_target_label = table.lookup(input_ops.features[target_name])
      outputs[PG_TARGET] = tf.squeeze(input_target_label)

    # TODO(brandondutra): get the score of the target label too.
    probabilities = model_fn_ops.predictions['probabilities']

    # if top_n == 0, this means use all the classes. We will use class names as
    # probabilities labels.
    if args.top_n == 0:
      predicted_index = tf.argmax(probabilities, axis=1)
      predicted = table.lookup(predicted_index)
      outputs.update({PG_CLASSIFICATION_FIRST_LABEL: predicted})
      probabilities_list = tf.unstack(probabilities, axis=1)
      for class_name, p in zip(class_names, probabilities_list):
        outputs[class_name] = p
    else:
      top_n = args.top_n

      # get top k labels and their scores.
      (top_k_values, top_k_indices) = tf.nn.top_k(probabilities, k=top_n)
      top_k_labels = table.lookup(tf.to_int64(top_k_indices))

      # Write the top_k values using 2*top_n columns.
      num_digits = int(math.ceil(math.log(top_n, 10)))
      if num_digits == 0:
        num_digits = 1
      for i in range(0, top_n):
        # Pad i based on the size of k. So if k = 100, i = 23 -> i = '023'. This
        # makes sorting the columns easy.
        padded_i = str(i + 1).zfill(num_digits)

        if i == 0:
          label_alias = PG_CLASSIFICATION_FIRST_LABEL
        else:
          label_alias = PG_CLASSIFICATION_LABEL_TEMPLATE % padded_i

        label_tensor_name = (tf.squeeze(
            tf.slice(top_k_labels, [0, i], [tf.shape(top_k_labels)[0], 1])))

        if i == 0:
          score_alias = PG_CLASSIFICATION_FIRST_SCORE
        else:
          score_alias = PG_CLASSIFICATION_SCORE_TEMPLATE % padded_i

        score_tensor_name = (tf.squeeze(
            tf.slice(top_k_values,
                     [0, i],
                     [tf.shape(top_k_values)[0], 1])))

        outputs.update({label_alias: label_tensor_name,
                        score_alias: score_tensor_name})

  else:
    if keep_target:
      outputs[PG_TARGET] = tf.squeeze(input_ops.features[target_name])

    scores = model_fn_ops.predictions['scores']
    outputs[PG_REGRESSION_PREDICTED_TARGET] = tf.squeeze(scores)

  return outputs

# This function is strongly based on
# tensorflow/contrib/learn/python/learn/estimators/estimator.py:export_savedmodel()
# The difference is we need to modify estimator's output layer.
def make_export_strategy(
        args,
        keep_target,
        assets_extra,
        features,
        schema,
        stats):
  """Makes prediction graph that takes json input.

  Args:
    args: command line args
    keep_target: If ture, target column is returned in prediction graph. Target
        column must also exist in input data
    assets_extra: other fiels to copy to the output folder
    job_dir: root job folder
    features: features dict
    schema: schema list
    stats: stats dict
  """
  target_name = feature_transforms.get_target_name(features)
  csv_header = [col['name'] for col in schema]
  if not keep_target:
    csv_header.remove(target_name)

  def export_fn(estimator, export_dir_base, checkpoint_path=None, eval_result=None):
    with ops.Graph().as_default() as g:
      contrib_variables.create_global_step(g)

      input_ops = feature_transforms.build_csv_serving_tensors_for_training_step(
          args.analysis, features, schema, stats, keep_target)
      model_fn_ops = estimator._call_model_fn(input_ops.features,
                                              None,
                                              model_fn_lib.ModeKeys.INFER)
      output_fetch_tensors = make_prediction_output_tensors(
          args=args,
          features=features,
          input_ops=input_ops,
          model_fn_ops=model_fn_ops,
          keep_target=keep_target)

      # Don't use signature_def_utils.predict_signature_def as that renames
      # tensor names if there is only 1 input/output tensor!
      signature_inputs = {key: tf.saved_model.utils.build_tensor_info(tensor)
                          for key, tensor in six.iteritems(input_ops.default_inputs)}
      signature_outputs = {key: tf.saved_model.utils.build_tensor_info(tensor)
                           for key, tensor in six.iteritems(output_fetch_tensors)}
      signature_def_map = {
          'serving_default':
              signature_def_utils.build_signature_def(
                  signature_inputs,
                  signature_outputs,
                  tf.saved_model.signature_constants.PREDICT_METHOD_NAME)}

      if not checkpoint_path:
        # Locate the latest checkpoint
        checkpoint_path = saver.latest_checkpoint(estimator._model_dir)
      if not checkpoint_path:
        raise ValueError("Couldn't find trained model at %s."
                         % estimator._model_dir)

      export_dir = saved_model_export_utils.get_timestamped_export_dir(
          export_dir_base)

      if (model_fn_ops.scaffold is not None and
         model_fn_ops.scaffold.saver is not None):
        saver_for_restore = model_fn_ops.scaffold.saver
      else:
        saver_for_restore = saver.Saver(sharded=True)

      with tf_session.Session('') as session:
        saver_for_restore.restore(session, checkpoint_path)
        init_op = control_flow_ops.group(
            variables.local_variables_initializer(),
            resources.initialize_resources(resources.shared_resources()),
            tf.tables_initializer())

        # Perform the export
        builder = saved_model_builder.SavedModelBuilder(export_dir)
        builder.add_meta_graph_and_variables(
            session, [tag_constants.SERVING],
            signature_def_map=signature_def_map,
            assets_collection=ops.get_collection(
                ops.GraphKeys.ASSET_FILEPATHS),
            legacy_init_op=init_op)
        builder.save(False)

      # Add the extra assets
      if assets_extra:
        assets_extra_path = os.path.join(compat.as_bytes(export_dir),
                                         compat.as_bytes('assets.extra'))
        for dest_relative, source in assets_extra.items():
          dest_absolute = os.path.join(compat.as_bytes(assets_extra_path),
                                       compat.as_bytes(dest_relative))
          dest_path = os.path.dirname(dest_absolute)
          file_io.recursive_create_dir(dest_path)
          file_io.copy(source, dest_absolute)

    # only keep the last 3 models
    saved_model_export_utils.garbage_collect_exports(
        export_dir_base,
        exports_to_keep=3)

    # save the last model to the model folder.
    # export_dir_base = A/B/intermediate_models/
    if keep_target:
      final_dir = os.path.join(args.job_dir, 'evaluation_model')
    else:
      final_dir = os.path.join(args.job_dir, 'model')
    if file_io.is_directory(final_dir):
      file_io.delete_recursively(final_dir)
    file_io.recursive_create_dir(final_dir)
    recursive_copy(export_dir, final_dir)

    return export_dir

  if keep_target:
    intermediate_dir = 'intermediate_evaluation_models'
  else:
    intermediate_dir = 'intermediate_prediction_models'

  return export_strategy.ExportStrategy(intermediate_dir, export_fn)

def get_estimator(args, output_dir, features, stats, target_vocab_size):
  # Check layers used for dnn models.
  if is_dnn_model(args.model) and not args.hidden_layer_sizes:
    raise ValueError('--hidden-layer-size* must be used with DNN models')
  if is_linear_model(args.model) and args.hidden_layer_sizes:
    raise ValueError('--hidden-layer-size* cannot be used with linear models')

  # Build tf.learn features
  feature_columns = build_feature_columns(features, stats, args.model)

  # Set how often to run checkpointing in terms of steps.
  config = tf.contrib.learn.RunConfig(
      save_checkpoints_steps=args.min_eval_frequency)

  train_dir = os.path.join(output_dir, 'train')
  if args.model == 'dnn_regression':
    estimator = tf.contrib.learn.DNNRegressor(
        feature_columns=feature_columns,
        hidden_units=args.hidden_layer_sizes,
        config=config,
        model_dir=train_dir,
        optimizer=tf.train.AdamOptimizer(
            args.learning_rate, epsilon=args.epsilon))
  elif args.model == 'linear_regression':
    estimator = tf.contrib.learn.LinearRegressor(
        feature_columns=feature_columns,
        config=config,
        model_dir=train_dir,
        optimizer=tf.train.FtrlOptimizer(
            args.learning_rate,
            l1_regularization_strength=args.l1_regularization,
            l2_regularization_strength=args.l2_regularization))
  elif args.model == 'dnn_classification':
    estimator = tf.contrib.learn.DNNClassifier(
        feature_columns=feature_columns,
        hidden_units=args.hidden_layer_sizes,
        n_classes=target_vocab_size,
        config=config,
        model_dir=train_dir,
        optimizer=tf.train.AdamOptimizer(
            args.learning_rate, epsilon=args.epsilon))
  elif args.model == 'linear_classification':
    estimator = tf.contrib.learn.LinearClassifier(
        feature_columns=feature_columns,
        n_classes=target_vocab_size,
        config=config,
        model_dir=train_dir,
        optimizer=tf.train.FtrlOptimizer(
            args.learning_rate,
            l1_regularization_strength=args.l1_regularization,
            l2_regularization_strength=args.l2_regularization))
  else:
    raise ValueError('bad --model-type value')

  return estimator

def read_vocab(args, column_name):
  """Reads a vocab file if it exists.

  Args:
    args: command line flags
    column_name: name of column to that has a vocab file.

  Returns:
    List of vocab words or [] if the vocab file is not found.
  """
  vocab_path = os.path.join(args.analysis,
                            feature_transforms.VOCAB_ANALYSIS_FILE % column_name)

  if not file_io.file_exists(vocab_path):
    return []

  vocab, _ = feature_transforms.read_vocab_file(vocab_path)
  return vocab

def get_key_names(features):
  names = []
  for name, transform in six.iteritems(features):
    if transform['transform'] == feature_transforms.KEY_TRANSFORM:
      names.append(name)
  return names

def read_json_file(file_path):
  if not file_io.file_exists(file_path):
    raise ValueError('File not found: %s' % file_path)
  return json.loads(file_io.read_file_to_string(file_path).decode())

def get_experiment_fn(args):
  """Builds the experiment function for learn_runner.run.

  Args:
    args: the command line args

  Returns:
    A function that returns a tf.learn experiment object.
  """

  def get_experiment(output_dir):
    # Read schema, input features, and transforms.
    schema_path_with_target = os.path.join(args.analysis,
                                           feature_transforms.SCHEMA_FILE)
    features_path = os.path.join(args.analysis,
                                 feature_transforms.FEATURES_FILE)
    stats_path = os.path.join(args.analysis,
                              feature_transforms.STATS_FILE)

    schema = read_json_file(schema_path_with_target)
    features = read_json_file(features_path)
    stats = read_json_file(stats_path)

    target_column_name = feature_transforms.get_target_name(features)
    if not target_column_name:
      raise ValueError('target missing from features file.')

    # Make a copy of the schema file without the target column.
    schema_without_target = [col for col in schema if col['name'] != target_column_name]
    schema_path_without_target = os.path.join(args.job_dir, 'schema_without_target.json')
    file_io.recursive_create_dir(args.job_dir)
    file_io.write_string_to_file(schema_path_without_target,
                                 json.dumps(schema_without_target, indent=2))

    # Make list of files to save with the trained model.
    additional_assets_with_target = {
        feature_transforms.FEATURES_FILE: features_path,
        feature_transforms.SCHEMA_FILE: schema_path_with_target}
    additional_assets_without_target = {
        feature_transforms.FEATURES_FILE: features_path,
        feature_transforms.SCHEMA_FILE: schema_path_without_target}

    # Get the model to train.
    target_vocab = read_vocab(args, target_column_name)
    estimator = get_estimator(args, output_dir, features, stats, len(target_vocab))

    export_strategy_csv_notarget = make_export_strategy(
        args=args,
        keep_target=False,
        assets_extra=additional_assets_without_target,
        features=features,
        schema=schema,
        stats=stats)
    export_strategy_csv_target = make_export_strategy(
        args=args,
        keep_target=True,
        assets_extra=additional_assets_with_target,
        features=features,
        schema=schema,
        stats=stats)

    # Build readers for training.
    if args.transform:
      if any(v['transform'] == feature_transforms.IMAGE_TRANSFORM
             for k, v in six.iteritems(features)):
        raise ValueError('"image_to_vec" transform requires transformation step. ' +
                         'Cannot train from raw data.')

      input_reader_for_train = feature_transforms.build_csv_transforming_training_input_fn(
          schema=schema,
          features=features,
          stats=stats,
          analysis_output_dir=args.analysis,
          raw_data_file_pattern=args.train,
          training_batch_size=args.train_batch_size,
          num_epochs=args.max_epochs,
          randomize_input=True,
          min_after_dequeue=10,
          reader_num_threads=multiprocessing.cpu_count())
      input_reader_for_eval = feature_transforms.build_csv_transforming_training_input_fn(
          schema=schema,
          features=features,
          stats=stats,
          analysis_output_dir=args.analysis,
          raw_data_file_pattern=args.eval,
          training_batch_size=args.eval_batch_size,
          num_epochs=1,
          randomize_input=False,
          reader_num_threads=multiprocessing.cpu_count())
    else:
      input_reader_for_train = feature_transforms.build_tfexample_transfored_training_input_fn(
          schema=schema,
          features=features,
          analysis_output_dir=args.analysis,
          raw_data_file_pattern=args.train,
          training_batch_size=args.train_batch_size,
          num_epochs=args.max_epochs,
          randomize_input=True,
          min_after_dequeue=10,
          reader_num_threads=multiprocessing.cpu_count())
      input_reader_for_eval = feature_transforms.build_tfexample_transfored_training_input_fn(
          schema=schema,
          features=features,
          analysis_output_dir=args.analysis,
          raw_data_file_pattern=args.eval,
          training_batch_size=args.eval_batch_size,
          num_epochs=1,
          randomize_input=False,
          reader_num_threads=multiprocessing.cpu_count())

    if args.early_stopping_num_evals == 0:
      train_monitors = None
    else:
      if is_classification_model(args.model):
        early_stop_monitor = tf.contrib.learn.monitors.ValidationMonitor(
            input_fn=input_reader_for_eval,
            every_n_steps=args.min_eval_frequency,
            early_stopping_rounds=(args.early_stopping_num_evals * args.min_eval_frequency),
            early_stopping_metric='accuracy',
            early_stopping_metric_minimize=False)
      else:
        early_stop_monitor = tf.contrib.learn.monitors.ValidationMonitor(
            input_fn=input_reader_for_eval,
            every_n_steps=args.min_eval_frequency,
            early_stopping_rounds=(args.early_stopping_num_evals * args.min_eval_frequency))
      train_monitors = [early_stop_monitor]

    return tf.contrib.learn.Experiment(
        estimator=estimator,
        train_input_fn=input_reader_for_train,
        eval_input_fn=input_reader_for_eval,
        train_steps=args.max_steps,
        train_monitors=train_monitors,
        export_strategies=[export_strategy_csv_notarget, export_strategy_csv_target],
        min_eval_frequency=args.min_eval_frequency,
        eval_steps=None)

  # Return a function to create an Experiment.
  return get_experiment

def local_analysis(args):
  if args.analysis:
    # Already analyzed.
    return

  if not args.schema or not args.features:
    raise ValueError('Either --analysis, or both --schema and --features are provided.')

  tf_config = json.loads(os.environ.get('TF_CONFIG', '{}'))
  cluster_spec = tf_config.get('cluster', {})
  if len(cluster_spec.get('worker', [])) > 0:
    raise ValueError('If "schema" and "features" are provided, local analysis will run and ' +
                     'only BASIC scale-tier (no workers node) is supported.')

  if cluster_spec and not (args.schema.startswith('gs://') and args.features.startswith('gs://')):
    raise ValueError('Cloud trainer requires GCS paths for --schema and --features.')

  print('Running analysis.')
  schema = json.loads(file_io.read_file_to_string(args.schema).decode())
  features = json.loads(file_io.read_file_to_string(args.features).decode())
  args.analysis = os.path.join(args.job_dir, 'analysis')
  args.transform = True
  file_io.recursive_create_dir(args.analysis)
  feature_analysis.run_local_analysis(args.analysis, args.train, schema, features)
  print('Analysis done.')

def set_logging_level(args):
  if 'TF_CONFIG' in os.environ:
    tf.logging.set_verbosity(tf.logging.INFO)
  else:
    tf.logging.set_verbosity(tf.logging.ERROR)
  if args.logging_level == 'error':
    tf.logging.set_verbosity(tf.logging.ERROR)
  elif args.logging_level == 'warning':
    tf.logging.set_verbosity(tf.logging.WARN)
  elif args.logging_level == 'info':
    tf.logging.set_verbosity(tf.logging.INFO)

def main(argv=None):
  args = parse_arguments(sys.argv if argv is None else argv)
  local_analysis(args)
  set_logging_level(args)
  # Supress TensorFlow Debugging info.
  os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

  learn_runner.run(
      experiment_fn=get_experiment_fn(args),
      output_dir=args.job_dir)

if __name__ == '__main__':
  main()


import pygame
import colorsys

SWATCH_WIDTH = 6
RECT_SIZE = 80
TEXT_HEIGHT = 15

def set_text(text):
    """Set font and text"""
    my_font = pygame.font.SysFont("Helvetica", RECT_SIZE // 6)
    text_surface = my_font.render(text, False, (255, 255, 255))
    return text_surface

def x_offset(count):
    x = count % SWATCH_WIDTH * RECT_SIZE
    return x

def y_offset(count):
    y = count // SWATCH_WIDTH * RECT_SIZE
    return y

def start_pygame(colors):
    # Pygame junk
    pygame.init()
    # Prepare Screen
    # screensize = (min(SWATCH_WIDTH, len(colors)) * RECT_SIZE, (1 + len(colors) // SWATCH_WIDTH) * RECT_SIZE)
    screensize = (min(SWATCH_WIDTH, len(colors)) * RECT_SIZE, 6 * (RECT_SIZE + 15))

    # Set caption and clock rate
    pygame.display.set_caption("Color Palettes")
    clock = pygame.time.Clock()
    surface = pygame.display.set_mode(screensize)
    fps = 1

    # Game loop
    run = True
    while run:
        clock.tick(fps)
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                run = False
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_RETURN:
                    run = False
                elif event.key == pygame.K_g:
                    pass
        draw(colors, surface)
        pygame.display.update()
    pygame.quit()

def draw(colors, surface):
    surface.blit(set_text("Original Color:"), (0, y_offset(0)))
    for count, color in enumerate(colors[0:1]):
        (r, g, b) = colorsys.hsv_to_rgb(color.h / 360.0, color.s / 100.0, color.v / 100.0)
        pygame.draw.rect(surface, (r * 255, g * 255, b * 255),
                         pygame.Rect(x_offset(count), y_offset(count) + 15, RECT_SIZE, RECT_SIZE))
    # Complementary color
    surface.blit(set_text("Complementary Color:"), (0, y_offset(0) + (RECT_SIZE + TEXT_HEIGHT)))
    for count, color in enumerate(colors[0:2]):
        (r, g, b) = colorsys.hsv_to_rgb(color.h / 360.0, color.s / 100.0, color.v / 100.0)
        pygame.draw.rect(surface, (r * 255, g * 255, b * 255),
                         pygame.Rect(x_offset(count), y_offset(count) + RECT_SIZE + 30, RECT_SIZE, RECT_SIZE))
    # Analogous colors
    surface.blit(set_text("Analogous Colors:"), (0, y_offset(0) + 2 * (RECT_SIZE + TEXT_HEIGHT)))
    for count, color in enumerate(colors[2:5]):
        (r, g, b) = colorsys.hsv_to_rgb(color.h / 360.0, color.s / 100.0, color.v / 100.0)
        pygame.draw.rect(surface, (r * 255, g * 255, b * 255),
                         pygame.Rect(x_offset(count), y_offset(count) + 2 * (RECT_SIZE + 15) + 15, RECT_SIZE,
                                     RECT_SIZE))
    # Split complementary colors
    surface.blit(set_text("Split Complementary Colors:"), (0, y_offset(0) + 3 * (RECT_SIZE + TEXT_HEIGHT)))
    for count, color in enumerate(colors[5:8]):
        (r, g, b) = colorsys.hsv_to_rgb(color.h / 360.0, color.s / 100.0, color.v / 100.0)
        pygame.draw.rect(surface, (r * 255, g * 255, b * 255),
                         pygame.Rect(x_offset(count), y_offset(count) + 3 * (RECT_SIZE + 15) + 15, RECT_SIZE,
                                     RECT_SIZE))
    # Triad colors
    surface.blit(set_text("Triad Colors:"), (0, y_offset(count) + 4 * (RECT_SIZE + TEXT_HEIGHT)))
    for count, color in enumerate(colors[8:11]):
        (r, g, b) = colorsys.hsv_to_rgb(color.h / 360.0, color.s / 100.0, color.v / 100.0)
        pygame.draw.rect(surface, (r * 255, g * 255, b * 255),
                         pygame.Rect(x_offset(count), y_offset(count) + 4 * (RECT_SIZE + 15) + 15, RECT_SIZE,
                                     RECT_SIZE))
    # Tetradic colors
    surface.blit(set_text("Tetradic Colors:"), (0, y_offset(count) + 5 * (RECT_SIZE + TEXT_HEIGHT)))
    for count, color in enumerate(colors[11:15]):
        (r, g, b) = colorsys.hsv_to_rgb(color.h / 360.0, color.s / 100.0, color.v / 100.0)
        pygame.draw.rect(surface, (r * 255, g * 255, b * 255),
                         pygame.Rect(x_offset(count), y_offset(count) + 5 * (RECT_SIZE + 15) + 15, RECT_SIZE,
                                     RECT_SIZE))