package net.minecraft.world.entity.ai.behavior;

import java.util.function.Supplier;
import java.util.Comparator;
import java.util.stream.Stream;
import java.util.Map;
import com.google.common.collect.ImmutableMap;
import com.mojang.datafixers.types.DynamicOps;
import java.util.function.Function;
import com.mojang.datafixers.Dynamic;
import com.google.common.collect.Lists;
import java.util.Random;
import java.util.List;

public class WeightedList<U> {
    protected final List<WeightedEntry<? extends U>> entries;
    private final Random random;
    
    public WeightedList(final Random random) {
        this.entries = Lists.newArrayList();
        this.random = random;
    }
    
    public WeightedList() {
        this(new Random());
    }
    
    public <T> WeightedList(final Dynamic<T> dynamic, final Function<Dynamic<T>, U> function) {
        this();
        final U object5;
        final int integer6;
        dynamic.asStream().forEach(dynamic -> dynamic.get("data").map(dynamic3 -> {
            object5 = function.apply(dynamic3);
            integer6 = dynamic.get("weight").asInt(1);
            return this.add(object5, integer6);
        }));
    }
    
    public <T> T serialize(final DynamicOps<T> dynamicOps, final Function<U, Dynamic<T>> function) {
        return (T)dynamicOps.createList((Stream)this.streamEntries().map(a -> dynamicOps.createMap((Map)ImmutableMap.builder().put(dynamicOps.createString("data"), function.apply(a.getData()).getValue()).put(dynamicOps.createString("weight"), dynamicOps.createInt(a.getWeight())).build())));
    }
    
    public WeightedList<U> add(final U object, final int integer) {
        this.entries.add(new WeightedEntry<U>(object, integer));
        return this;
    }
    
    public WeightedList<U> shuffle() {
        return this.shuffle(this.random);
    }
    
    public WeightedList<U> shuffle(final Random random) {
        this.entries.forEach(a -> a.setRandom(random.nextFloat()));
        this.entries.sort(Comparator.comparingDouble(object -> object.getRandWeight()));
        return this;
    }
    
    public Stream<? extends U> stream() {
        return this.entries.stream().map(WeightedEntry::getData);
    }
    
    public Stream<WeightedEntry<? extends U>> streamEntries() {
        return this.entries.stream();
    }
    
    public U getOne(final Random random) {
        return (U)this.shuffle(random).stream().findFirst().<Throwable>orElseThrow(RuntimeException::new);
    }
    
    @Override
    public String toString() {
        return "WeightedList[" + this.entries + "]";
    }
    
    public class WeightedEntry<T> {
        private final T data;
        private final int weight;
        private double randWeight;
        
        private WeightedEntry(final T object, final int integer) {
            this.weight = integer;
            this.data = object;
        }
        
        private double getRandWeight() {
            return this.randWeight;
        }
        
        private void setRandom(final float float1) {
            this.randWeight = -Math.pow(float1, 1.0f / this.weight);
        }
        
        public T getData() {
            return this.data;
        }
        
        public int getWeight() {
            return this.weight;
        }
        
        @Override
        public String toString() {
            return "" + this.weight + ":" + this.data;
        }
    }
}
