#%% Libraries

#% Libraries: Custom

#%%
class ClusterGroup(object):
    def __init__(self, cluster_name):
        self.cluster_name = cluster_name
        self.clusters = {}
        self.connects = {
            'input': [],
            'output': []
        }
        self.back_channels = {
        }
        
    def parse(self):
        snapshot = {
            'type': type(self).__name__,
            'cluster_name': self.cluster_name,
            'connects': self.connects
        }
        snapshot['clusters'] = [cluster.parse() for cluster in self.clusters.values()]
        return snapshot
    
    def grouped_cluster_name(self, cluster_name):
        return '%s:%s' % (self.cluster_name, cluster_name)
    
    def change_cluster_name(self, new_cluster_name):
        for key, value in list(self.clusters.items()):
            changed_cluster_name = '%s:%s' % (new_cluster_name, value.cluster_name.split(':')[-1])
            value.change_cluster_name(changed_cluster_name)
            self.clusters[changed_cluster_name] = self.clusters.pop(key)
            
        for key, value in self.connects.items():
            self.connects[key] = ['%s:%s' % (new_cluster_name, cluster_name.split(':')[-1]) for cluster_name in value]
            
        self.cluster_name = new_cluster_name
        
    def add_backchannel(self, backchannel_name, inputs = [], outputs = []):
        if backchannel_name in self.back_channels.keys():
            raise Exception('backchannel %s already exists' % backchannel_name)
            
        self.back_channels[backchannel_name] = {
            'input': inputs,
            'output': outputs
        }
        
    def add_back_connect(self, backchannel_name, connect_type, cluster_name):
        if backchannel_name not in self.back_channels.keys():
            raise Exception('backchannel %s not in cluster group' % backchannel_name)
        
        if cluster_name not in self.clusters.keys():
            raise Exception('cluster %s not in cluster group' % cluster_name)
            
        self.back_channels[backchannel_name][connect_type].append(cluster_name)
        
    def coef_count(self):
        coefs = sum([cluster.coef_count() for cluster in self.clusters.values()])
        return coefs
    
    ## Mimmic Controller Function
    def add_cluster(self, cluster_name, cluster_class, *args, **kwargs):
        grouped_cluster_name = self.grouped_cluster_name(cluster_name)
        if grouped_cluster_name in self.clusters.keys():
            raise Exception('cluster_name %s already exists' % cluster_name)
            
        self.clusters[grouped_cluster_name] = cluster_class(grouped_cluster_name, *args, **kwargs)
        
    def set_cluster_connect(self, connect_type, cluster_name):
        grouped_cluster_name = self.grouped_cluster_name(cluster_name)
        self.connects[connect_type].append(grouped_cluster_name)
        
    def get_cluster(self, cluster_name):
        grouped_cluster_name = self.grouped_cluster_name(cluster_name)
        return self.clusters[grouped_cluster_name]
    
    def connect_clusters(self, from_cluster, to_cluster, *args):
        self.get_cluster(from_cluster).connect_add('output', self.get_cluster(to_cluster), *args)
        self.get_cluster(to_cluster).connect_add('input', self.get_cluster(from_cluster), *args)
        
    def connect_clusters_many(self, connect_items):
        for connect_item in connect_items:
            self.connect_clusters(*connect_item)
            
    def add_with_connects(self, cluster_name, cluster_class, connect_items, *args, **kwargs):
        
        def gen_new_connects(connect_items):
            connect_list = []
            for connect_item in connect_items:
                direction = connect_item[0]
                connect_name = connect_item[1]
                
                if direction == 'input':
                    connect_list.append((connect_name, cluster_name, *connect_item[2:]))
                elif direction == 'output':
                    connect_list.append((cluster_name, connect_name, *connect_item[2:]))
                    
            return connect_list
        
        self.add_cluster(cluster_name, cluster_class, *args, **kwargs)
        self.connect_clusters_many(gen_new_connects(connect_items))
            
    ## Mimmic Cluster Function
    def return_cluster(self, connect_type):
        return [self.clusters[cluster_name] for cluster_name in self.connects[connect_type]]
    
    def init_cluster(self, obs_count, train_index):
        for cluster in self.clusters.values():
            cluster.init_cluster(obs_count, train_index)
#%%
