package alien4cloud.tosca.parser.postprocess;

import java.util.Collections;
import java.util.Map;
import java.util.Set;

import org.alien4cloud.tosca.model.templates.Capability;
import org.alien4cloud.tosca.model.templates.NodeTemplate;
import org.alien4cloud.tosca.model.types.AbstractToscaType;
import org.alien4cloud.tosca.model.types.CapabilityType;
import org.springframework.stereotype.Component;

import com.google.common.collect.Maps;
import com.google.common.collect.Sets;

import alien4cloud.tosca.context.ToscaContext;
import lombok.Setter;

@Component
public class CapabilityMatcherService implements ICapabilityMatcherService {

    public interface IToscaContextFinder {
        <T extends AbstractToscaType> T find(Class<T> clazz, String id);
    }

    @Setter
    private IToscaContextFinder toscaContextFinder;

    public CapabilityMatcherService() {
        toscaContextFinder = ToscaContext::get;
    }

    @Override
    public Map<String, Capability> getCompatibleCapabilityByType(NodeTemplate nodeTemplate, String type) {
        Map<String, Capability> capabilities = nodeTemplate.getCapabilities();
        if (capabilities == null) {
            return Collections.emptyMap();
        }

        Map<String, Capability> targetCapabilitiesMatch = Maps.newHashMap();
        Set<String> allTypes = Sets.newHashSet();
        for (Map.Entry<String, Capability> capabilityEntry : capabilities.entrySet()) {
            CapabilityType capabilityType = toscaContextFinder.find(CapabilityType.class, capabilityEntry.getKey());
            allTypes.clear();
            allTypes.add(capabilityType.getElementId());
            allTypes.addAll(capabilityType.getDerivedFrom());

            if (allTypes.contains(type)) {
                targetCapabilitiesMatch.put(capabilityEntry.getKey(), capabilityEntry.getValue());
            }
        }
        return targetCapabilitiesMatch;
    }
}

@@ -95,30 +95,8 @@
 Map;
-%0Aimport java.util.Set;
 %0A%0Aim
@@ -322,24 +322,83 @@
 bilityType;%0A
+import org.alien4cloud.tosca.normative.ToscaNormativeUtil;%0A
 import org.s
@@ -477,47 +477,8 @@
 aps;
-%0Aimport com.google.common.collect.Sets;
 %0A%0Aim
@@ -1288,58 +1288,8 @@
 ();%0A
-        Set%3CString%3E allTypes = Sets.newHashSet();%0A
     
@@ -1388,65 +1388,16 @@
     
-CapabilityType capabilityType = toscaContextFinder.find(C
+String c
 apab
@@ -1409,15 +1409,14 @@
 Type
-.class,
+Name =
  cap
@@ -1436,17 +1436,16 @@
 getKey()
-)
 ;%0A      
@@ -1454,52 +1454,65 @@
     
-allTypes.clear();%0A            allTypes.add(c
+CapabilityType capabilityType = toscaContextFinder.find(C
 apab
@@ -1525,123 +1525,100 @@
 ype.
-getElementId());%0A            allTypes.addAll(capabilityType.getDerivedFrom());%0A%0A            if (allTypes.contains(t
+class, capabilityTypeName);%0A%0A            if (ToscaNormativeUtil.isFromType(type, capabilityT
 ype)
@@ -1675,30 +1675,24 @@
 pability
-Entry.getKey()
+TypeName
 , capabi

Simplify code