package org.qcri.rheem.core.optimizer.cardinality;

import org.qcri.rheem.core.api.Configuration;
import org.qcri.rheem.core.optimizer.OptimizationContext;
import org.qcri.rheem.core.plan.executionplan.Channel;
import org.qcri.rheem.core.plan.rheemplan.*;
import org.qcri.rheem.core.platform.ExecutionState;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Collection;
import java.util.Collections;

/**
 * Handles the {@link CardinalityEstimate}s of a {@link RheemPlan}.
 */
public class CardinalityEstimatorManager {

    private final Logger logger = LoggerFactory.getLogger(this.getClass());

    /**
     * The {@link RheemPlan} whose cardinalities are being managed.
     */
    private final RheemPlan rheemPlan;

    /**
     * Keeps the {@link CardinalityEstimate}s around.
     */
    private final OptimizationContext optimizationContext;

    /**
     * Provides {@link CardinalityEstimator}s etc.
     */
    private final Configuration configuration;

    private CardinalityEstimationTraversal planTraversal;

    public CardinalityEstimatorManager(RheemPlan rheemPlan,
                                       OptimizationContext optimizationContext,
                                       Configuration configuration) {
        this.rheemPlan = rheemPlan;
        this.optimizationContext = optimizationContext;
        this.configuration = configuration;
    }

    public void pushCardinalities() {
        this.getPlanTraversal().traverse(this.optimizationContext, this.configuration);
        this.optimizationContext.clearMarks();
        assert this.optimizationContext.isTimeEstimatesComplete();
    }

    public CardinalityEstimationTraversal getPlanTraversal() {
        if (this.planTraversal == null) {
            this.planTraversal = CardinalityEstimationTraversal.createPushTraversal(
                    Collections.emptyList(),
                    this.rheemPlan.collectReachableTopLevelSources(),
                    this.configuration
            );
        }
        return this.planTraversal;
    }

    /**
     * Injects the cardinalities of a current {@link ExecutionState} into its associated {@link RheemPlan}
     * (or its {@link OptimizationContext}, respectively) and then reperforms the cardinality estimation.
     */
    public void pushCardinalityUpdates(ExecutionState executionState) {
        this.injectMeasuredCardinalities(executionState);
        this.pushCardinalities();
    }

    /**
     * Injects the cardinalities of a current {@link ExecutionState} into its associated {@link RheemPlan}.
     */
    private void injectMeasuredCardinalities(ExecutionState executionState) {
        executionState.getCardinalityMeasurements().forEach(this::injectMeasureCardinality);
    }

    /**
     * Injects the measured {@code cardinality} of a {@code channel} into the {@link #optimizationContext}.
     */
    private void injectMeasureCardinality(Channel channel, long cardinality) {
        // Build the new CardinalityEstimate.
        final CardinalityEstimate newEstimate = new CardinalityEstimate(cardinality, cardinality, 1d);

        // Identify the Slots that correspond to the channel.
        final Collection<Slot<?>> correspondingSlots = channel.getCorrespondingSlots();
        for (Slot<?> correspondingSlot : correspondingSlots) {

            // Identify the corresponding OperatorContext.
            final Operator owner = correspondingSlot.getOwner();
            final OptimizationContext.OperatorContext operatorCtx = this.optimizationContext.getOperatorContext(owner);
            if (operatorCtx == null) {
                // FIXME: Within loops, we might need to propagate across iterations!
                this.logger.debug("Could not inject measured cardinality for {}: It is presumably a glue operator or inside of a loop.", owner);
                continue;
            }

            // Update the operatorCtx, then propagate.
            // TODO: Do we need to propagate? Revisit once channel.getCorrespondingSlots() is incorporating OptimizationContexts.
            final int slotIndex = correspondingSlot.getIndex();
            if (correspondingSlot instanceof InputSlot<?>) {
                operatorCtx.setInputCardinality(slotIndex, newEstimate);
                owner.propagateInputCardinality(slotIndex, operatorCtx);
            } else {
                assert correspondingSlot instanceof OutputSlot<?>;
                operatorCtx.setOutputCardinality(slotIndex, newEstimate);
                owner.propagateOutputCardinality(slotIndex, operatorCtx);
            }
        }
    }

}

@@ -3464,64 +3464,12 @@
     
-final OptimizationContext.OperatorContext operatorCtx = 
+if (
 this
@@ -3518,37 +3518,8 @@
 ner)
-;%0A            if (operatorCtx
  == 
@@ -3869,126 +3869,122 @@
     
-// TODO: Do we need to propagate? Revisit once channel.getCorrespondingSlots() is incorporating OptimizationContexts.%0A
+if (correspondingSlot instanceof InputSlot%3C?%3E) %7B%0A                // Find the outermost InputSlot and propagate.%0A  
     
@@ -3995,120 +3995,306 @@
     
+  
 final 
-int slotIndex = correspondingSlot.getIndex();%0A            if (correspondingSlot instanceof InputSlot%3C?%3E) %7B
+InputSlot%3C?%3E outerInput = owner.getOutermostInputSlot((InputSlot%3C?%3E) correspondingSlot);%0A                final Operator outerOperator = outerInput.getOwner();%0A                final OptimizationContext.OperatorContext operatorCtx = this.optimizationContext.getOperatorContext(outerOperator);
 %0A   
@@ -4334,33 +4334,45 @@
 Cardinality(
-slo
+outerInput.ge
 tIndex
+()
 , newEstimat
@@ -4384,35 +4384,43 @@
                o
-wne
+uterOperato
 r.propagateInput
@@ -4427,33 +4427,45 @@
 Cardinality(
-slo
+outerInput.ge
 tIndex
+()
 , operatorCt
@@ -4460,32 +4460,33 @@
 , operatorCtx);%0A
+%0A
             %7D el
@@ -4490,16 +4490,80 @@
  else %7B%0A
+                // Find the outermost OutputSlot and propagate.%0A
         
@@ -4621,16 +4621,497 @@
 lot%3C?%3E;%0A
+                @SuppressWarnings(%22unchecked%22)%0A                final Collection%3COutputSlot%3C?%3E%3E outerOutputs = owner.getOutermostOutputSlots((OutputSlot) correspondingSlot);%0A                for (OutputSlot%3C?%3E outerOutput : outerOutputs) %7B%0A                    final Operator outerOperator = outerOutput.getOwner();%0A                    final OptimizationContext.OperatorContext operatorCtx =%0A                            this.optimizationContext.getOperatorContext(outerOperator);%0A    
         
@@ -5151,25 +5151,38 @@
 inality(
-slo
+outerOutput.ge
 tIndex
+()
 , newEst
@@ -5205,20 +5205,32 @@
         
+  
-owne
+  outerOperato
 r.propag
@@ -5254,17 +5254,30 @@
 ity(
-slo
+outerOutput.ge
 tIndex
+()
 , op
@@ -5288,16 +5288,34 @@
 orCtx);%0A
+                %7D%0A
         

correct cardinality injection before re-optimization