using System;
using System.Collections.Generic;
using System.Linq;
using Tensorflow.Common.Types;
using Tensorflow.Keras.Engine;
using Tensorflow.Train;
using Tensorflow.Training;
using Tensorflow.Training.Saving.SavedModel;
using static Tensorflow.Binding;

namespace Tensorflow.Hub
{
    public class KerasLayer : Layer
    {
        private string _handle;
        private LoadOptions? _load_options;
        private Trackable _func;
        private Func _callable;

        public KerasLayer(string handle, bool trainable = false, LoadOptions? load_options = null) :
            base(new Keras.ArgsDefinition.LayerArgs() { Trainable = trainable })
        {
            _handle = handle;
            _load_options = load_options;

            _func = load_module(_handle, _load_options);
            _track_trackable(_func, "_func");
            // TODO(Rinne): deal with _is_hub_module_v1.

            _callable = _get_callable();
            _setup_layer(trainable);
        }

        private void _setup_layer(bool trainable = false)
        {
            HashSet trainable_variables;
            if (_func is Layer layer)
            {
                foreach (var v in layer.TrainableVariables)
                {
                    _add_existing_weight(v, true);
                }
                trainable_variables = new HashSet(layer.TrainableVariables.Select(v => v.UniqueId));
            }
            else if (_func.CustomizedFields.TryGetValue("trainable_variables", out var obj) && obj is IEnumerable trackables)
            {
                foreach (var trackable in trackables)
                {
                    if (trackable is IVariableV1 v)
                    {
                        _add_existing_weight(v, true);
                    }
                }
                trainable_variables = new HashSet(trackables.Where(t => t is IVariableV1).Select(t => (t as IVariableV1).UniqueId));
            }
            else
            {
                trainable_variables = new HashSet();
            }

            if (_func is Layer)
            {
                layer = (Layer)_func;
                foreach (var v in layer.Variables)
                {
                    if (!trainable_variables.Contains(v.UniqueId))
                    {
                        _add_existing_weight(v, false);
                    }
                }
            }
            else if (_func.CustomizedFields.TryGetValue("variables", out var obj) && obj is IEnumerable total_trackables)
            {
                foreach (var trackable in total_trackables)
                {
                    if (trackable is IVariableV1 v && !trainable_variables.Contains(v.UniqueId))
                    {
                        _add_existing_weight(v, false);
                    }
                }
            }

            if (_func.CustomizedFields.ContainsKey("regularization_losses"))
            {
                if ((_func.CustomizedFields["regularization_losses"] as ListWrapper)?.Count > 0)
                {
                    throw new NotImplementedException("The regularization_losses loading has not been supported yet, " +
                        "please submit an issue to https://github.com/SciSharp/TensorFlow.NET/issues to let us know and add a feature.");
                }
            }
        }

        protected override Tensors Call(Tensors inputs, Tensors state = null, bool? training = null, IOptionalArgs? optionalArgs = null)
        {
            _check_trainability();

            // TODO(Rinne): deal with training_argument

            var result = _callable(inputs);

            return _apply_output_shape_if_set(inputs, result);
        }

        private void _check_trainability()
        {
            if (!Trainable) return;

            // TODO(Rinne): deal with _is_hub_module_v1 and signature

            if (TrainableWeights is null || TrainableWeights.Count == 0)
            {
                tf.Logger.Error("hub.KerasLayer is trainable but has zero trainable weights.");
            }
        }

        private Tensors _apply_output_shape_if_set(Tensors inputs, Tensors result)
        {
            // TODO(Rinne): implement it.
            return result;
        }

        private void _add_existing_weight(IVariableV1 weight, bool? trainable = null)
        {
            bool is_trainable;
            if (trainable is null)
            {
                is_trainable = weight.Trainable;
            }
            else
            {
                is_trainable = trainable.Value;
            }
            add_weight(weight.Name, weight.shape, weight.dtype, trainable: is_trainable, getter: x => weight);
        }

        private Func _get_callable()
        {
            if (_func is Layer layer)
            {
                return x => layer.Apply(x);
            }
            if (_func.CustomizedFields.ContainsKey("__call__"))
            {
                if (_func.CustomizedFields["__call__"] is RestoredFunction function)
                {
                    return x => function.Apply(x);
                }
            }
            throw new ValueError("Cannot get the callable from the model.");
        }

        private static Trackable load_module(string handle, LoadOptions? load_options = null)
        {
            //var set_load_options = load_options ?? LoadContext.get_load_option();
            return module_v2.load(handle, load_options);
        }
    }
}