package com.conveyal.r5.profile;

import com.conveyal.r5.transit.TransitLayer;
import com.conveyal.r5.transit.TripPattern;
import com.conveyal.r5.transit.TripSchedule;
import gnu.trove.iterator.TIntObjectIterator;
import gnu.trove.map.TIntObjectMap;
import gnu.trove.map.hash.TIntObjectHashMap;
import org.apache.commons.math3.random.MersenneTwister;

import java.util.Arrays;

/**
 * Stores random offsets for frequency trips.
 * This is not in RaptorWorkerData as RaptorWorkerData may be shared between threads.
 */
public class FrequencyRandomOffsets {
    /** map from trip pattern index to a list of offsets for trip i and frequency entry j on that pattern */
    public final TIntObjectMap<int[][]> offsets = new TIntObjectHashMap<>();
    public final TransitLayer data;

    /** The mersenne twister is a higher quality random number generator than the one included with Java */
    private MersenneTwister mt = new MersenneTwister();

    public FrequencyRandomOffsets(TransitLayer data) {
        this.data = data;

        if (!data.hasFrequencies)
            return;

        for (int pattIdx = 0; pattIdx < data.tripPatterns.size(); pattIdx++) {
            TripPattern tp = data.tripPatterns.get(pattIdx);

            if (!tp.hasFrequencies) continue;

            int[][] offsetsThisPattern = new int[tp.tripSchedules.size()][];

            for (int tripIdx = 0; tripIdx < tp.tripSchedules.size(); tripIdx++) {
                TripSchedule ts = tp.tripSchedules.get(tripIdx);
                offsetsThisPattern[tripIdx] = ts.headwaySeconds == null ? null : new int[ts.headwaySeconds.length];
            }

            offsets.put(pattIdx, offsetsThisPattern);
        }
    }

    public void randomize () {
        int remaining = 0;

        // first clear old offsets
        for (TIntObjectIterator<int[][]> it = offsets.iterator(); it.hasNext(); ) {
            it.advance();
            for (int[] offsetsPerEntry : it.value()) {
                Arrays.fill(offsetsPerEntry, -1);
                remaining += offsetsPerEntry.length;
            }
        }

        while (remaining > 0) {
            int remainingAfterPreviousRound = remaining;

            for (TIntObjectIterator<int[][]> it = offsets.iterator(); it.hasNext(); ) {
                it.advance();

                TripPattern pattern = data.tripPatterns.get(it.key());
                int[][] val = it.value();

                for (int tripScheduleIndex = 0; tripScheduleIndex < val.length; tripScheduleIndex++) {
                    TripSchedule schedule = pattern.tripSchedules.get(tripScheduleIndex);

                    // it is possible to have both frequency and non-frequency trips on the same pattern. If this is a scehduled
                    // trip, don't set offset
                    if (schedule.headwaySeconds == null)
                        val[tripScheduleIndex] = null;
                    else {
                        for (int frequencyEntryIndex = 0; frequencyEntryIndex < val[tripScheduleIndex].length; frequencyEntryIndex++) {
                            if (schedule.phaseFromId == null || schedule.phaseFromId[frequencyEntryIndex] == null) {
                                // not phased. also, don't overwrite with new random number on each iteration, as other
                                // trips may be phased from this one
                                if (val[tripScheduleIndex][frequencyEntryIndex] == -1) {
                                    val[tripScheduleIndex][frequencyEntryIndex] = mt.nextInt(schedule.headwaySeconds[frequencyEntryIndex]);
                                    remaining--;
                                }
                            }
                            else {
                                if (val[tripScheduleIndex][frequencyEntryIndex] != -1) continue; // already randomized

                                // find source phase information
                                int[] source = data.frequencyEntryIndexForId.get(schedule.phaseFromId[frequencyEntryIndex]);

                                int sourcePatternIdx = source[0];
                                int sourceTripScheduleIdx = source[1];
                                int sourceFrequencyEntryIdx = source[2];

                                // check if it's already been randomized
                                int previousOffset = offsets.get(sourcePatternIdx)[sourceTripScheduleIdx][sourceFrequencyEntryIdx];

                                if (previousOffset != -1) {
                                    TripPattern phaseFromPattern = data.tripPatterns.get(sourcePatternIdx);
                                    TripSchedule phaseFromSchedule = phaseFromPattern.tripSchedules.get(sourceTripScheduleIdx);

                                    // figure out stop indices
                                    int sourceStopIndexInPattern = 0;
                                    int sourceStopIndexInNetwork = data.indexForStopId.get(schedule.phaseFromStop[frequencyEntryIndex]);

                                    // TODO check that stop IDs were found.

                                    while (sourceStopIndexInPattern < phaseFromPattern.stops.length &&
                                            phaseFromPattern.stops[sourceStopIndexInPattern] != sourceStopIndexInNetwork) {
                                        sourceStopIndexInPattern++;
                                    }

                                    if (sourceStopIndexInPattern == phaseFromPattern.stops.length) {
                                        throw new IllegalArgumentException(String.format("Stop %s was not found in source pattern!", schedule.phaseFromStop[frequencyEntryIndex]));
                                    }

                                    int targetStopIndexInPattern = 0;
                                    int targetStopIndexInNetwork = data.indexForStopId.get(schedule.phaseAtStop[frequencyEntryIndex]);

                                    while (targetStopIndexInPattern < pattern.stops.length &&
                                            pattern.stops[targetStopIndexInPattern] != targetStopIndexInNetwork) {
                                        targetStopIndexInPattern++;
                                    }

                                    // TODO This should really be checked also before modifications are applied.
                                    if (targetStopIndexInPattern == pattern.stops.length) {
                                        throw new IllegalArgumentException(String.format("Stop %s was not found in target pattern!", schedule.phaseAtStop[frequencyEntryIndex]));
                                    }

                                    // use arrivals at last stop
                                    int[] sourceTravelTimes = sourceStopIndexInPattern < phaseFromPattern.stops.length - 1 ?
                                            phaseFromSchedule.departures : phaseFromSchedule.arrivals;

                                    // figure out the offset if they were to pass the stops at the same time
                                    int timeAtSourceStop = phaseFromSchedule.startTimes[sourceFrequencyEntryIdx] +
                                            sourceTravelTimes[sourceStopIndexInPattern] +
                                            previousOffset;

                                    // use arrivals at last stop
                                    int[] targetTravelTimes = targetStopIndexInPattern < pattern.stops.length - 1 ?
                                            schedule.departures : schedule.arrivals;

                                    // figure out when the target trip passes the stop if the offset were 0.
                                    int timeAtTargetStop = schedule.startTimes[frequencyEntryIndex] +
                                            targetTravelTimes[targetStopIndexInPattern];

                                    int offset = timeAtSourceStop - timeAtTargetStop;

                                    // this is the offset so the trips arrive at the same time. We now add the desired phase.
                                    offset += schedule.phaseSeconds[frequencyEntryIndex];

                                    // make sure it's positive
                                    while (offset < 0) offset += schedule.headwaySeconds[frequencyEntryIndex];

                                    // make it as small as possible...
                                    offset %= schedule.headwaySeconds[frequencyEntryIndex];

                                    val[tripScheduleIndex][frequencyEntryIndex] = offset;
                                    remaining--;
                                }
                            }
                        }
                    }
                }
            }

            if (remainingAfterPreviousRound == remaining && remaining > 0) {
                throw new IllegalArgumentException("Cannot solve phasing, you may have a circular reference!");
            }
        }
    }
}

@@ -4028,32 +4028,565 @@
 cyEntryIndex%5D);%0A
+                                // Throw a meaningful error when invalid IDs are encountered instead of NPE.%0A                                // Really this should be done when applying the modifications rather than during the search.%0A                                if (source == null) %7B%0A                                    throw new RuntimeException(%22This pattern ID specified in a scenario does not exist: %22%0A                                            + schedule.phaseFromId%5BfrequencyEntryIndex%5D);%0A                                %7D
 %0A               

throw meaningful exception