from VLMP.components.models import modelBase

import copy

class SIMULATION(modelBase):
    """
    Component name: SIMULATION
    Component type: model

    Author: Pablo Ibáñez-Freire
    Date: 6/11/2023

    Load model from dictionary generated by pyUAMMD simulation

    :param inputSimulation: Input dictionary generated by pyUAMMD simulation
    :type inputSimulation: dict
    """

    def __init__(self,name,**params):
        super().__init__(_type = self.__class__.__name__,
                         _name= name,
                         availableParameters = {"inputSimulation"},
                         requiredParameters  = {"inputSimulation"},
                         definedSelections   = {"particleId","forceField"},
                         **params)

        ############################################################

        inputSimulation = params["inputSimulation"]
        self.logger.info(f"[SIMULATION] Created model from pyUAMMD simulation")

        ########################################################

        # TYPES

        types = self.getTypes()

        typesLabels = inputSimulation["global"]["types"]["labels"]
        for typ in inputSimulation["global"]["types"]["data"]:
            typInfo = {l:typ[i] for i,l in enumerate(typesLabels)}
            types.addType(**typInfo)

        #Generate positions, a line along the z axis
        state = copy.deepcopy(inputSimulation["state"])

        #Generate structure
        structure = copy.deepcopy(inputSimulation["topology"]["structure"])

        #Generate forceField
        forceField = copy.deepcopy(inputSimulation["topology"]["forceField"])

        ########################################################

        self.setState(state)
        self.setStructure(structure)
        self.setForceField(forceField)


    def processSelection(self,**params):

        sel = []

        if "particleId" in params:
            sel += params["particleId"]

        if "forceField" in params:
            sel += self.getForceFieldSelection(params["forceField"])

        return sel
