package io.openems.edge.energy.optimizer;

import static com.google.common.collect.ImmutableMap.toImmutableMap;
import static com.google.common.collect.ImmutableSet.toImmutableSet;
import static com.google.common.collect.ImmutableSortedMap.toImmutableSortedMap;

import java.time.ZonedDateTime;
import java.time.format.DateTimeFormatter;
import java.util.Comparator;
import java.util.HashMap;
import java.util.Locale;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Optional;
import java.util.TreeMap;
import java.util.stream.IntStream;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.ImmutableSortedMap;

import io.jenetics.Genotype;
import io.openems.edge.energy.api.handler.DifferentModes;
import io.openems.edge.energy.api.handler.DifferentModes.Period.Transition;
import io.openems.edge.energy.api.handler.EnergyScheduleHandler;
import io.openems.edge.energy.api.handler.Fitness;
import io.openems.edge.energy.api.simulation.EnergyFlow;
import io.openems.edge.energy.api.simulation.GlobalOptimizationContext;
import io.openems.edge.energy.api.simulation.GlobalOptimizationContext.Ess;
import io.openems.edge.energy.api.simulation.GlobalOptimizationContext.Period.Hour;
import io.openems.edge.energy.api.simulation.GlobalOptimizationContext.Period.Quarter;
import io.openems.edge.energy.api.simulation.GocUtils;
import io.openems.edge.energy.api.simulation.periods.PeriodData.Price;
import io.openems.edge.energy.api.simulation.periods.Periods;
import io.openems.edge.energy.optimizer.ModeCombinations.ModeCombination;

public record SimulationResult(//
		Fitness fitness, //
		ImmutableSortedMap<ZonedDateTime, Period> periods, //
		ImmutableMap<//
				? extends EnergyScheduleHandler.WithDifferentModes, //
				ImmutableSortedMap<ZonedDateTime, DifferentModes.Period.Transition>> schedules, //
		ImmutableSet<? extends EnergyScheduleHandler.WithOnlyOneMode> eshsWithOnlyOneMode, //
		Ess ess, //
		int simulationsCounter, //
		int generationsCounter) {

	/**
	 * A Period in a {@link SimulationResult}. Duration of one period is always one
	 * quarter.
	 */
	public record Period(//
			GlobalOptimizationContext.Period period, //
			ModeCombination modeCombination, //
			EnergyFlow energyFlow, //
			int essInitialEnergy //
	) {

		/**
		 * Constructor for {@link Period}.
		 * 
		 * @param gocPeriod        the {@link GlobalOptimizationContext.Period}
		 * @param energyFlow       the {@link EnergyFlow}
		 * @param essInitialEnergy the initial ESS energy in the beginning of the period
		 *                         in [Wh]
		 * @return a {@link Period}
		 */
		public static Period from(GlobalOptimizationContext.Period gocPeriod, ModeCombination modeCombination,
				EnergyFlow energyFlow, int essInitialEnergy) {
			return new Period(gocPeriod, modeCombination, energyFlow, essInitialEnergy);
		}
	}

	/**
	 * An empty {@link SimulationResult}.
	 */
	public static final SimulationResult EMPTY_SIMULATION_RESULT = new SimulationResult(Fitness.builder().build(), //
			ImmutableSortedMap.of(), ImmutableMap.of(), ImmutableSet.of(), new Ess(0, 0, 0, 0), 0, 0);

	protected static class BestScheduleCollector {
		private final ImmutableSortedMap.Builder<ZonedDateTime, SimulationResult.Period> periods = //
				ImmutableSortedMap.naturalOrder();
		private final Map<EnergyScheduleHandler.WithDifferentModes, //
				TreeMap<ZonedDateTime, Integer>> modesPerEsh = new HashMap<>();

		protected void addPeriod(ZonedDateTime time, SimulationResult.Period period) {
			this.periods.put(time, period);
		}

		protected synchronized void addMode(ZonedDateTime time, EnergyScheduleHandler.WithDifferentModes esh,
				int postProcessedMode) {
			var modes = this.modesPerEsh.computeIfAbsent(esh, e -> new TreeMap<ZonedDateTime, Integer>());
			modes.put(time, postProcessedMode);
		}
	}

	/**
	 * Re-Simulate a {@link Genotype} to create a {@link SimulationResult}.
	 *
	 * @param goc                the {@link GlobalOptimizationContext}
	 * @param schedule           the schedule as defined by {@link EshCodec}
	 * @param simulationsCounter the total number of simulations
	 * @param generationsCounter the total number of generations
	 * @return the {@link SimulationResult}
	 */
	private static SimulationResult from(//
			GlobalOptimizationContext goc, //
			int[] schedule, //
			int simulationsCounter, //
			int generationsCounter) {
		final var bsc = new BestScheduleCollector();
		final var fitness = Simulator.simulate(goc, ModeCombinations.fromGlobalOptimizationContext(goc), schedule, bsc,
				GocUtils.normalizeEshModePreferenceRanks(goc.eshsWithDifferentModes()));
		final var periods = bsc.periods.buildOrThrow();

		var schedules = bsc.modesPerEsh.entrySet().stream() //
				.collect(toImmutableMap(Entry::getKey, // ESH
						e1 -> {
							ImmutableSortedMap<ZonedDateTime, Transition> subMap = e1.getValue().entrySet().stream() //
									.collect(toImmutableSortedMap(//
											Comparator.naturalOrder(), //
											Entry::getKey, // time
											e2 -> { // Period.Transition
												final var mode = e2.getValue();
												final var p = periods.get(e2.getKey());
												final var gridBuyPrice = p.period.data().gridBuyPrice()//
														.map(Price::actual)//
														.orElse(null);
												final var gridSellPrice = p.period.data().gridSellPrice()//
														.map(Price::actual)//
														.orElse(null);

												return new DifferentModes.Period.Transition(p.period.duration(), mode, //
														gridBuyPrice, gridSellPrice, p.energyFlow(),
														p.essInitialEnergy());
											}));
							return subMap;
						}));

		var eshsWithOnlyOneMode = goc.eshs().stream() //
				.filter(EnergyScheduleHandler.WithOnlyOneMode.class::isInstance)
				.map(EnergyScheduleHandler.WithOnlyOneMode.class::cast) //
				.collect(toImmutableSet());

		return new SimulationResult(//
				fitness.build(), //
				periods, //
				schedules, //
				eshsWithOnlyOneMode, //
				goc.ess(), //
				simulationsCounter, //
				generationsCounter);
	}

	/**
	 * Re-Simulate a {@link Genotype} to create a {@link SimulationResult}.
	 * 
	 * <p>
	 * This method re-simulates using the {@link Quarter} periods and not (only) the
	 * {@link Hour} periods.
	 * 
	 * @param goc                the {@link GlobalOptimizationContext}
	 * @param schedule           the schedule as defined by {@link EshCodec}
	 * @param simulationsCounter the total number of simulations
	 * @param generationsCounter the total number of generations
	 * @return the {@link SimulationResult}
	 */
	public static SimulationResult fromQuarters(//
			GlobalOptimizationContext goc, //
			int[] schedule, //
			int simulationsCounter, //
			int generationsCounter) {
		if (goc == null || schedule.length == 0) {
			return EMPTY_SIMULATION_RESULT;
		}

		// Convert to Quarters
		final var quarterPeriods = Periods.copyOfQuarterly(goc.periods());
		final var quarterGoc = new GlobalOptimizationContext(goc.clock(), goc.environment(), goc.startTime(),
				goc.eshs(), goc.eshsWithDifferentModes(), goc.grid(), goc.ess(), quarterPeriods);
		final var quarterSchedule = IntStream.range(0, goc.periods().size()) //
				.flatMap(periodIndex -> switch (goc.periods().get(periodIndex)) {
				case GlobalOptimizationContext.Period.Hour ph //
					-> ph.quarterPeriods().stream().mapToInt(ignore -> periodIndex); // repeat
				case GlobalOptimizationContext.Period.Quarter pq //
					-> IntStream.of(periodIndex);
				}) //
				.map(periodIndex -> periodIndex < schedule.length //
						? schedule[periodIndex] //
						: 0) // fallback
				// TODO use default index
				.toArray();

		return from(quarterGoc, quarterSchedule, simulationsCounter, generationsCounter);
	}

	private static final DateTimeFormatter TIME_FORMATTER = DateTimeFormatter.ofPattern("HH:mm");

	private static void log(StringBuilder b, String format, Object... args) {
		b.append(String.format(Locale.ENGLISH, format, args).toString());
	}

	/**
	 * Builds a log string of this {@link SimulationResult}.
	 * 
	 * @return log string
	 */
	public String toLogString(String prefix) {
		var b = new StringBuilder(prefix);
		var firstValue = Optional.ofNullable(this.periods.firstEntry()) //
				.map(Entry::getValue) //
				.orElse(null);
		if (firstValue == null) {
			return b.append("NO PERIODS").toString();
		}

		b.append("Time BuyLimit GridBuyPrice GridSellPrice Prod  Cons MCons   Ess  Grid  EssInitial");
		firstValue.energyFlow.getManagedConsumptions().keySet() //
				.forEach(v -> log(b, " %-10s", v.substring(Math.max(0, v.length() - 10))));
		b.append("\n");

		this.periods.entrySet().forEach(e -> {
			final var time = e.getKey();
			final var p = e.getValue();
			final var ef = p.energyFlow;
			log(b, "%s", prefix);
			log(b, "%s ", time.format(TIME_FORMATTER));
			Optional.ofNullable(p.period.gridBuySoftLimit()).ifPresentOrElse(//
					limit -> log(b, "%7d ", limit), //
					() -> log(b, "%7s ", "-"));
			p.period.data().gridBuyPrice()//
					.map(Price::actual)//
					.ifPresentOrElse(//
							v -> log(b, "%5.0f ", v), //
							() -> log(b, "            -"));
			p.period.data().gridSellPrice()//
					.map(Price::actual)//
					.ifPresentOrElse(//
							v -> log(b, "%5.0f ", v), //
							() -> log(b, "            -"));
			log(b, "%5d ", ef.getProduction());
			log(b, "%5d ", ef.getUnmanagedConsumption());
			log(b, "%5d ", ef.getConsumption());
			log(b, "%5d ", ef.getEss());
			log(b, "%5d ", ef.getGrid());
			log(b, "%10d ", p.essInitialEnergy);
			ef.getManagedConsumptions().values().stream() //
					.forEach(v -> log(b, "%10d", v));
			this.schedules.forEach((esh, schedule) -> {
				log(b, " %-10s ", esh.modes().getAsString(schedule.get(time).modeIndex()));
			});
			b.append("\n");
		});
		b.append(prefix)//
				.append("totalNumberOfSimulations=").append(this.simulationsCounter)//
				.append(";totalNumberOfGenerations=").append(this.generationsCounter)//
				.append(";fitness=").append(this.fitness);
		return b.toString();
	}
}
