/*
 *  ******************************************************************************
 *  *
 *  *
 *  * This program and the accompanying materials are made available under the
 *  * terms of the Apache License, Version 2.0 which is available at
 *  * https://www.apache.org/licenses/LICENSE-2.0.
 *  *
 *  *  See the NOTICE file distributed with this work for additional
 *  *  information regarding copyright ownership.
 *  * Unless required by applicable law or agreed to in writing, software
 *  * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 *  * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 *  * License for the specific language governing permissions and limitations
 *  * under the License.
 *  *
 *  * SPDX-License-Identifier: Apache-2.0
 *  *****************************************************************************
 */
package org.deeplearning4j.rl4j.network;

import lombok.NonNull;
import org.deeplearning4j.nn.graph.ComputationGraph;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;

public class NetworkHelper {

    /**
     * Create a {@link INetworkHandler} with explicit network inputs to channel names mapping
     * @param model A {@link ComputationGraph} instance
     * @param networkInputsToChannelNameMap A {@link ChannelToNetworkInputMapper.NetworkInputToChannelBinding} array
     *                                      describing which observation channels to feed to which network inputs.
     * @param channelNames The list of channel names generated by the transform process. Should be in the same order.
     * @param labelNames The names of the network's output labels
     * @param gradientName The name of the network's gradient
     * @return A {@link INetworkHandler}
     */
    public INetworkHandler buildHandler(ComputationGraph model,
                                        @NonNull ChannelToNetworkInputMapper.NetworkInputToChannelBinding[] networkInputsToChannelNameMap,
                                        String[] channelNames,
                                        String[] labelNames,
                                        String gradientName) {
        String[] networkInputNames = model.getConfiguration().getNetworkInputs().toArray(new String[0]);

        ChannelToNetworkInputMapper mapper = new ChannelToNetworkInputMapper(networkInputsToChannelNameMap, networkInputNames, channelNames);
        return new ComputationGraphHandler(model, labelNames, gradientName, mapper);
    }

    /**
     * Create a {@link INetworkHandler} with the single network input mapped to a specific observation channel
     * @param model A {@link ComputationGraph} instance
     * @param networkInputChannel The name of the observation channel to use as the network's input.
     *                            Empty to use the first channel.
     * @param channelNames The list of channel names generated by the transform process. Should be in the same order.
     * @param labelNames The names of the network's output labels
     * @param gradientName The name of the network's gradient
     * @return A {@link INetworkHandler}
     */
    public INetworkHandler buildHandler(ComputationGraph model,
                                        String networkInputChannel,
                                        String[] channelNames,
                                        String[] labelNames,
                                        String gradientName) {
        int channelIdx = findChannelIdx(channelNames, networkInputChannel);
        return new ComputationGraphHandler(model, labelNames, gradientName, channelIdx);
    }

    /**
     * Create a {@link INetworkHandler} with the single network input mapped to a specific observation channel
     * @param model A {@link MultiLayerNetwork} instance
     * @param networkInputChannel The name of the observation channel to use as the network's input.
     *                            Empty to use the first channel.
     * @param channelNames The list of channel names generated by the transform process. Should be in the same order.
     * @param labelName The name of the network's output label
     * @param gradientName The name of the network's gradient
     * @return A {@link INetworkHandler}
     */
    public INetworkHandler buildHandler(MultiLayerNetwork model,
                                        String networkInputChannel,
                                        String[] channelNames,
                                        String labelName,
                                        String gradientName) {
        int channelIdx = findChannelIdx(channelNames, networkInputChannel);
        return new MultiLayerNetworkHandler(model, labelName, gradientName, channelIdx);
    }

    private int findChannelIdx(String[] channelNames, String channelName) {
        // When the channel name or the channelNames is null or empty, always use the first channel
        if(channelName == null || channelName.isEmpty() || channelNames == null || channelNames.length == 0) {
            return 0;
        }

        for (int i = 0; i < channelNames.length; ++i) {
            if (channelNames[i] == channelName) {
                return i;
            }
        }

        throw new IllegalArgumentException("The channel '" + channelName + "' was not found in channelNames.");
    }
}
