package pl.edu.pw.mini.jk.sg.mixeduct.attacker;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import org.apache.commons.collections4.queue.CircularFifoQueue;
import org.apache.commons.math3.distribution.EnumeratedDistribution;
import org.apache.commons.math3.random.RandomGenerator;
import org.apache.commons.math3.util.Pair;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.Lists;
import org.slf4j.LoggerFactory;

import pl.edu.pw.mini.jk.sg.mixeduct.IterationResult;
import pl.edu.pw.mini.jk.sg.mixeduct.attacker.historychanger.AttackerHistoryChanger;

public class ProbabilisticStrategyAttackerWithHistoryFactory<M> implements AttackerWithHistory<M> {
    protected CircularFifoQueue<EnumeratedDistribution<ImmutableList<M>>> history;
    private final RandomGenerator rng;
    private final int initialHistoryLength;
    private final AttackerHistoryChanger ahc;    

    public ProbabilisticStrategyAttackerWithHistoryFactory(final ImmutableList<Pair<ImmutableList<M>, Double>> initialStrategy, int historyLength, final RandomGenerator rng, final AttackerHistoryChanger ahc) {
        this.history = new CircularFifoQueue<EnumeratedDistribution<ImmutableList<M>>>(historyLength);
        this.rng = rng;
	this.initialHistoryLength = historyLength;
	this.ahc = ahc;
	history.offer(new EnumeratedDistribution<ImmutableList<M>>(rng, initialStrategy));
    }

    @Override
    public AttackerStrategy<M> getInstance() {
	return new EqualProbableStrategyJoin<>(ImmutableList.copyOf(history), rng);
    }

    @Override
    public void registerIterationResult(final int iteration, final IterationResult<?, M, ?, ?, ?> iterationResult, final AttackerStrategy<M> trainingAttacker) {
	extendHistory(ahc.newHistoryLength(initialHistoryLength, history.size(), history.maxSize(), iteration, false));
        history.offer(new EnumeratedDistribution<>(rng, iterationResult.attackerStrategy().strategyAsPairs()));
	LoggerFactory.getLogger("petrus.AttackerHistory").debug("{} {}", iteration,
				   getCurrentHistory());

    }

    public void extendHistory(int newSize) {
        if (newSize != history.maxSize()) {
            CircularFifoQueue<EnumeratedDistribution<ImmutableList<M>>> newQueue = new CircularFifoQueue<>(
                   newSize);
            while(!history.isEmpty()){
                newQueue.add(history.poll());
            }
            history=newQueue;
        } else {

        }
    }
    
    public int getHistorySize() {
        return history.maxSize();
    }
    
    public int getCurrentHistory() {
    	return history.size();
    }

    public static<M> EnumeratedDistribution<ImmutableList<M>> joinDistributions(
            Collection<EnumeratedDistribution<ImmutableList<M>>> history, final double annealingCoefficient, final RandomGenerator rng) {
        final Map<ImmutableList<M>, Double> weights = new HashMap<>();
        double c = 1;

        List<EnumeratedDistribution<ImmutableList<M>>> reverse = Lists.reverse(new ArrayList<>(history));

        for (final EnumeratedDistribution<ImmutableList<M>> strategy : reverse) {
            final double cc = c;
            strategy.getPmf().stream().forEach(
                    p -> weights.put(p.getFirst(), weights.getOrDefault(p.getFirst(), 0.) + p.getSecond() * cc));
            c *= annealingCoefficient;
        }
	
        return new EnumeratedDistribution<ImmutableList<M>>(rng,
                weights.entrySet().stream().map(e -> new Pair<ImmutableList<M>, Double>(e.getKey(), e.getValue()))
                        .collect(Collectors.toList()));
    }
}
