# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/02_configs.ipynb (unless otherwise specified).

__all__ = ['BaseConfiguration', 'ProportionalConfiguration', 'ConfigurationStructure', 'DynamicConfigurationStructure',
           'DynamicConfiguration']

# Cell
import gym
import numpy as np
import random
from abc import ABC, abstractmethod

from pct.environments import PendulumV0
from pct.environments import VelocityModel
from pct.functions import IndexedParameter
from pct.functions import Constant
from pct.architectures import ProportionalArchitecture
from pct.architectures import DynamicArchitecture
from .structure import StructureDefinition
from pct.errors import BaseErrorCollector
from pct.architectures import LevelKey
from pct.environments import DummyModel
from epct.structure  import BinaryOnes
from pct.structure import ArchitectureStructure

# Cell
class BaseConfiguration(ABC):
    "Base class of an array configuration. This class is not used direclty by developers, but defines the functionality common to all."
    def __init__(self, name, num_inputs, num_actions, grid, references):
        self.name = name
        self.num_inputs=num_inputs
        self.num_actions=num_actions
        self.grid=grid
        self.references=references
        self.config={}
        self.config['parameters']={}
        if grid[-1] != len(references):
            raise ValueError(f'Number of references {len(references)} must equal number of top-level columns {grid[-1]}')

    def __call__(self):
        level0 = self.level0config()
        if len(self.grid) == 1:
            level0.append(self.references)
        self.config['level0']=level0

        if len(self.grid) > 2:
            for level in range(1,len(self.grid), 1):
                numColumnsThisLevel=self.grid[level]
                numColumnsPreviousLevel=self.grid[level-1]
                self.config[f'level{level}']=self.levelnconfig(numColumnsThisLevel, numColumnsPreviousLevel)

        if len(self.grid) > 1:
            self.config[f'level{len(self.grid)-1}'] = self.leveltopconfig(self.references, self.grid[-2])

        return self.config

    @classmethod
    def top_level(cls, config):
        levels = len(config.keys())
        if 'parameters' in config.keys():
            levels -= 1

        return levels-1

    def lists_of_ones(self, num_lists, length):
        wts=[]
        for i in range(num_lists):
            wt = [1 for iter in range(length)]

            wts.append(wt)

        return wts

    def create_binary_lists(self, num_lists, length):

        wts=[]
        for i in range(num_lists):
            if length==1 and num_lists==1:
                wt = [1]
            else:
                wt = [random.randint(0, 1) for iter in range(length)]
                if np.sum(wt) == 0:
                    index = random.randint(0, length-1)
                    wt[index] = 1

            wts.append(wt)


        #if length>1 and num_lists==1 and np.sum(wts) == 0:
        #    index = random.randint(0, length-1)
        #    wts[0][index] = 1

        # flip one list if all 0
        if length==1 and num_lists>1 and np.sum(wts) == 0:
            index = random.randint(0, num_lists-1)
            wts[index] = [1]

        return wts

    def create_randomint_lists(self, lower, upper, num_lists, length):
        wts=[]
        for i in range(num_lists):
            wt = [random.randint(lower, upper) for iter in range(length)]
            wts.append(wt)
        return wts

    def create_randomfloat_list(self, lower, upper, length):
        return [random.uniform(lower, upper) for iter in range(length)]

    def create_randomfloat_lists(self, lower, upper, num_lists, length):
        wts=[]
        for i in range(num_lists):
            wt = [random.uniform(lower, upper) for iter in range(length)]
            wts.append(wt)
        return wts

    def to_raw(self):
        raw = []
        for key in self.config.keys():
            if key != 'parameters':
                raw.append(self.config[key])
        return raw

    @classmethod
    def from_raw(cls, raw):
        config = {}
        config['parameters']={}
        for level in range(len(raw)):
            config[f'level{level}'] = raw[level]
        return config

    @classmethod
    def dict_to_raw(cls, config):
        raw = []
        for key in config.keys():
            if key != 'parameters':
                raw.append(config[key])
        return raw


    @abstractmethod
    def level0config(self):
         pass

    @abstractmethod
    def levelnconfig(self, numColumnsThisLevel, numColumnsPreviousLevel):
         pass

    @abstractmethod
    def leveltopconfig(self, top_references, numColumnsPreviousLevel):
         pass

    @classmethod
    def copy_data(cls, from_config, to_config):
        levels = len(from_config)-1
        for level in range(levels):
            index=0
            for item in range(len(from_config[f'level{level}'][index])):
                if item < len(to_config[f'level{level}'][index]):
                    for element in range(len(from_config[f'level{level}'][index][item])):
                        if element < len(to_config[f'level{level}'][index][item]):
                            to_config[f'level{level}'][index][item][element]= from_config[f'level{level}'][index][item][element]
                    BaseConfiguration.ensure_non_zero(to_config, level, index, item)

            index=1
            for item in range(len(from_config[f'level{level}'][index])):
                if item < len(to_config[f'level{level}'][index]):
                    to_config[f'level{level}'][index][item]= from_config[f'level{level}'][index][item]

            index=2
            for item in range(len(from_config[f'level{level}'][index])):
                #print(level, item)
                if item < len(to_config[f'level{level}'][index]):
                    for element in range(len(from_config[f'level{level}'][index][item])):
                        #print(level, item, element)
                        if element < len(to_config[f'level{level}'][index][item]):
                            #print(to_config[f'level{level}'][index][item][element])
                            #print(from_config[f'level{level}'][index][item][element])
                            to_config[f'level{level}'][index][item][element]= from_config[f'level{level}'][index][item][element]
                    BaseConfiguration.ensure_non_zero(to_config, level, index, item)


    @classmethod
    def copy_data_from_top(cls, from_config, to_config):
        levels = len(from_config)-1
        to_level = levels
        skip_level=levels-1

        for level in range(levels-1, -1, -1):
            if to_level==skip_level:
                to_level-=1
            #print(level, to_level, skip_level)
            index=0
            for item in range(len(from_config[f'level{level}'][index])):
                for element in range(len(from_config[f'level{level}'][index][item])):
                    to_config[f'level{to_level}'][index][item][element]= from_config[f'level{level}'][index][item][element]

            index=1
            for item in range(len(from_config[f'level{level}'][index])):
                to_config[f'level{to_level}'][index][item]= from_config[f'level{level}'][index][item]

            index=2
            for item in range(len(from_config[f'level{level}'][index])):
                for element in range(len(from_config[f'level{level}'][index][item])):
                    to_config[f'level{to_level}'][index][item][element]= from_config[f'level{level}'][index][item][element]
            to_level-=1



    @classmethod
    def copy_data_from_larger(cls, from_config, to_config):
        levels = len(from_config)-1
        to_level = levels-2
        skip_level=levels-2
        #print('skip_level',skip_level)
        for level in range(levels-1, -1, -1):
            if level==skip_level:
                pass #level-=1
            else:
                #print(level, to_level, skip_level)
                index=0
                for item in range(len(to_config[f'level{to_level}'][index])):
                    if to_level ==0 and len(to_config[f'level{to_level}'][index][item]) > 1:
                        for element in range(len(to_config[f'level{to_level}'][index][item])):
                            to_config[f'level{to_level}'][index][item][element]= from_config[f'level{level}'][index][item][element]

                index=1
                for item in range(len(to_config[f'level{to_level}'][index])):
                    to_config[f'level{to_level}'][index][item]= from_config[f'level{level}'][index][item]

                index=2
                for item in range(len(from_config[f'level{level}'][index])):
                    if item < len(to_config[f'level{to_level}'][index]):
                        for element in range(len(to_config[f'level{to_level}'][index][item])):
                            to_config[f'level{to_level}'][index][item][element]= from_config[f'level{level}'][index][item][element]

                to_level-=1
            if to_level < 0:
                break

    @classmethod
    def ensure_non_zero(cls, config, level, index, item):
        if np.sum(config[f'level{level}'][index][item])==0:
            length = len(config[f'level{level}'][index][item])
            ind = random.randint(0, length-1)
            config[f'level{level}'][index][item][ind]=1


    @classmethod
    def copy_data_from_smaller(cls, from_config, to_config):
        levels = len(from_config)-1
        to_level = 0
        skip_level=levels-1
        #print(to_config)
        for level in range(levels):
            if to_level==skip_level:
                to_level+=1
            #print(level, to_level, skip_level)
            index=0
            #print(len(from_config[f'level{level}'][index]), len(to_config[f'level{to_level}'][index]) )

            for item in range(len(from_config[f'level{level}'][index])):
                #print('item', item)
                #if item < len(to_config[f'level{to_level}'][index]) and item < len(from_config[f'level{level}'][index]):
                if item < len(to_config[f'level{to_level}'][index]) :
                    #print(len(from_config[f'level{level}'][index][item]), len(to_config[f'level{to_level}'][index][item]) )

                    for element in range(len(from_config[f'level{level}'][index][item])):
                        #print('elem', element)
                        if element < len(to_config[f'level{to_level}'][index][item]) and element < len(from_config[f'level{level}'][index][item]) :
                            to_config[f'level{to_level}'][index][item][element]= from_config[f'level{level}'][index][item][element]
                    BaseConfiguration.ensure_non_zero(to_config, to_level, index, item)
                #if np.sum(to_config[f'level{to_level}'][index][item])==0:
                #    length = len(to_config[f'level{to_level}'][index][item])
                #    ind = random.randint(0, length-1)
                #    to_config[f'level{to_level}'][index][item][ind]=1

            index=1
            for item in range(len(from_config[f'level{level}'][index])):
                to_config[f'level{to_level}'][index][item]= from_config[f'level{level}'][index][item]

            index=2
            for item in range(len(from_config[f'level{level}'][index])):
                if item < len(to_config[f'level{to_level}'][index]):
                    for element in range(len(from_config[f'level{level}'][index][item])):
                        to_config[f'level{to_level}'][index][item][element]= from_config[f'level{level}'][index][item][element]
            to_level+=1


# Cell
class ProportionalConfiguration(BaseConfiguration):
    "ProportionalConfiguration"
    def __init__(self, name="proportional", num_inputs=1, num_actions=1, grid=[2, 2], references=[1],
                 lower_float=-100, upper_float=100, **cargs):
        super().__init__(name, num_inputs, num_actions, grid, references)
        self.lower_float=lower_float
        self.upper_float=upper_float
        self.config['parameters']['lower_float']=lower_float
        self.config['parameters']['upper_float']=upper_float


    def level0config(self):
        numColumnsThisLevel = self.grid[0]
        # inputs
        iwts=self.create_binary_lists(self.num_inputs, numColumnsThisLevel)

        # outputs
        owts=self.create_randomfloat_list(self.lower_float, self.upper_float, numColumnsThisLevel)

        # actions
        #actwts = self.lists_of_ones(self.num_actions, numColumnsThisLevel)
        actwts = self.create_binary_lists(self.num_actions, numColumnsThisLevel)

        lwts=[]
        lwts.append(iwts)
        lwts.append(owts)
        lwts.append(actwts)

        return lwts

    def levelnconfig(self, numColumnsThisLevel, numColumnsPreviousLevel):
        # inputs
        iwts=self.create_binary_lists(numColumnsThisLevel, numColumnsPreviousLevel)

        # outputs
        owts=self.create_randomfloat_list(self.lower_float, self.upper_float, numColumnsThisLevel)

        # lower refs
        rwts = self.create_randomfloat_lists(self.lower_float, self.upper_float, numColumnsPreviousLevel, numColumnsThisLevel)

        lwts=[]
        lwts.append(iwts)
        lwts.append(owts)
        lwts.append(rwts)

        return lwts


    def leveltopconfig(self, top_references, numColumnsPreviousLevel):

        numColumnsThisLevel=len(top_references)
        # inputs
        iwts=self.create_binary_lists(numColumnsThisLevel, numColumnsPreviousLevel)

        # outputs
        owts=self.create_randomfloat_list(self.lower_float, self.upper_float, numColumnsThisLevel)

        # lower refs
        rwts = self.create_randomfloat_lists(self.lower_float, self.upper_float, numColumnsPreviousLevel, numColumnsThisLevel)

        lwts=[]
        lwts.append(iwts)
        lwts.append(owts)
        lwts.append(rwts)
        lwts.append(top_references)

        return lwts



# Cell
class ConfigurationStructure():

    def __init__(self, structure=None, **cargs):
        if structure==None:
            self.structure={}
        else:
            self.structure=structure

    def from_raw(self, raw):

        config = BaseConfiguration.from_raw(raw)
        levels = len(config)-1
        num_inputs=len(config['level0'][0])
        self.structure['num_inputs']=num_inputs
        num_actions=len(config['level0'][2])
        self.structure['num_actions']=num_actions

        grid=[]
        for level in range(levels):
            grid.append(len(config[f'level{level}'][1]))
        self.structure['grid']=grid
        references = config[f'level{levels-1}'][3]
        self.structure['references']=references
        return self.structure

    def add_nodes(self, level, num_nodes=1):
        mutated_structure = 0

        levels = len(self.structure['grid'])
        if level < levels-1:
            self.structure['grid'][level]+=num_nodes
            mutated_structure = 1

        return self.structure, mutated_structure

    def remove_nodes(self, level, num_nodes=1):
        mutated_structure = 0

        levels = len(self.structure['grid'])
        if level == levels-1:
            raise Exception("Cannot remove nodes from the top level.")

        if level > levels-1:
            raise Exception(f'Level {level} does not exist for hierarchy with levels {levels}.')

        existing_num_nodes = self.structure['grid'][level]
        if num_nodes > existing_num_nodes-1:
            raise Exception(f'Cannot remove {num_nodes} nodes from level with {existing_num_nodes} nodes.')

        self.structure['grid'][level]-=num_nodes
        mutated_structure = 1

        return self.structure, mutated_structure

    def add_level(self, num_columns=1):
        mutated_structure = 1

        #levels = len(self.structure['grid'])
        self.structure['grid'].insert(-1, num_columns)

        return self.structure, mutated_structure

    def remove_level(self):
        mutated_structure = 0

        levels = len(self.structure['grid'])
        if levels>1:
            del self.structure['grid'][levels-1:]
            self.structure['grid'][-1]=len(self.structure['references'])
            mutated_structure = 1

        return self.structure, mutated_structure


# Cell

class DynamicConfigurationStructure(ConfigurationStructure):

    def __init__(self, structure=None, **cargs):
        if structure==None:
            self.structure={}
        else:
            self.structure=structure

    def from_raw(self, raw):

        config = BaseConfiguration.from_raw(raw)
        levels = len(config)-1
        num_inputs=len(config['level0'][0])
        self.structure['num_inputs']=num_inputs
        num_actions=len(config['level0'][3])
        self.structure['num_actions']=num_actions

        grid=[]
        for level in range(levels):
            grid.append(len(config[f'level{level}'][1]))
        self.structure['grid']=grid
        references = config[f'level{levels-1}'][2]
        self.structure['references']=references
        return self.structure



# Cell
class DynamicConfiguration(BaseConfiguration):
    "DynamicConfiguration"
    def __init__(self, inputs=None, num_inputs=1, num_actions=1, references=None,
                 structure=None, grid=None,  **cargs):

        #self.inputs=inputs
        num_top_inputs=0

        top_inputs = structure.get_parameter('top_inputs')
        if top_inputs != None:
            num_top_inputs = len(top_inputs)
            if inputs != None:
                num_inputs = len(inputs)

        self.references=references
        if structure == None:
            structure = StructureDefinition(references=references)

        self.num_inputs=num_inputs - num_top_inputs
        self.num_top_inputs=num_top_inputs
        self.num_actions=num_actions
        self.config={}
        self.structure=structure
        self.grid=grid
        """
        self.lower_float=lower_float
        self.upper_float=upper_float
        self.config['parameters']['lower_float']=lower_float
        self.config['parameters']['upper_float']=upper_float
        """

    def __call__(self):

        if self.grid==None:
            grid=[]
            levels = random.randint(self.structure.get_parameter('min_levels_limit'), self.structure.get_parameter('levels_limit'))
            for level in range(levels-1):
                columns = random.randint(self.structure.get_parameter('min_columns_limit'), self.structure.get_parameter('columns_limit'))
                grid.append(columns)
            grid.append(len(self.references))
        else:
            grid=self.grid


        levels = len(grid)
        #print(grid)
        if levels == 1:
            level0 = self.level0topconfig(grid[0])
            self.config['level0']=level0
            return self.config
        else:
            level0 = self.level0config(grid[0], grid[1])

        self.config['level0']=level0

        if levels > 2:
            for level in range(1,levels-1, 1):
                numColumnsThisLevel=grid[level]
                numColumnsPreviousLevel=grid[level-1]
                numColumnsNextLevel = grid[level+1]
                self.config[f'level{level}']=self.levelnconfig(numColumnsThisLevel, numColumnsNextLevel, numColumnsPreviousLevel)

        if levels > 1:
            self.config[f'level{len(grid)-1}'] = self.leveltopconfig(grid[-1], grid[-2])

        return self.config

    def level0config(self, numColumnsThisLevel, numColumnsNextLevel):
        return self.structure.get_level0(self.num_inputs,numColumnsThisLevel, numColumnsNextLevel, self.num_actions)

    def level0topconfig(self, num_columns):
        return self.structure.get_level0top(self.num_inputs+self.num_top_inputs,num_columns, self.num_actions)

    def levelnconfig(self, numColumnsThisLevel, numColumnsNextLevel, numColumnsPreviousLevel):
        return self.structure.get_leveln(numColumnsThisLevel, numColumnsNextLevel, numColumnsPreviousLevel)

    def leveltopconfig(self, numColumnsThisLevel, numColumnsPreviousLevel):
        return self.structure.get_leveltop(numColumnsThisLevel, numColumnsPreviousLevel, self.num_top_inputs)

    @classmethod
    def convert_to_raw_from_proportional_raw(cls, prop_raw, verbose=False):
        if verbose:
            #print(prop_raw)
            for level in range(len(prop_raw)):
                if level == 0:
                    print(level, 'P', prop_raw[level][0],'O', prop_raw[level][1],'A', prop_raw[level][2])
                elif level == len(prop_raw)-1:
                    print(level, 'P', prop_raw[level][0],'O', prop_raw[level][1],'R', prop_raw[level][2],'MR', prop_raw[level][3])
                else:
                    print(level, 'P', prop_raw[level][0],'O', prop_raw[level][1],'R', prop_raw[level][2])


        levels = len(prop_raw)
        if levels == 0:
            dyn_raw=prop_raw
        else:
            dyn_raw=[]
            for level in range(len(prop_raw)):
                level_raw=[]
                if level == 0:
                    level_raw.append(prop_raw[level][0])
                else:
                    plist=[]
                    for item in range(len(prop_raw[level][0][0])):
                        item_list = []
                        for list in prop_raw[level][0]:
                            item_list.append(list[item])
                        plist.append(item_list)
                    level_raw.append(plist)
                level_raw.append(prop_raw[level][1])
                if level == 0:
                    level_raw.append(prop_raw[level+1][2])
                if level == levels-1:
                    level_raw.append(prop_raw[level][3])
                else:
                    level_raw.append(prop_raw[level][2])
                dyn_raw.append(level_raw)

        if verbose:
            print()

            for level in range(len(dyn_raw)):
                if level == 0:
                    print(level, 'P', dyn_raw[level][0],'O', dyn_raw[level][1],'R', dyn_raw[level][2],'A', dyn_raw[level][3])
                else:
                    print(level, 'P', dyn_raw[level][0],'O', dyn_raw[level][1],'R', dyn_raw[level][2])


        return dyn_raw

