#
# Copyright (C) 2013 - 2023 Oracle and/or its affiliates. All rights reserved.
#

from typing import List, Dict, Optional, Union

from pypgx._utils.error_handling import java_handler
from pypgx._utils import conversion
from pypgx.api.mllib._graphwise_conv_layer_config import (
    GraphWiseAttentionLayerConfig,
    GraphWiseConvLayerConfig,
)
from pypgx.api.mllib._edge_combination_method import EdgeCombinationMethod
from pypgx.api.mllib._input_property_config import InputPropertyConfig
from pypgx.api.mllib._mllib_conversion import (
    _java_input_property_list_to_python_configs,
    _java_conv_layer_configs_to_python_conv_layer_configs,
)
from pypgx.api.mllib._edge_combination_method import (
    ConcatEdgeCombinationMethod,
    ProductEdgeCombinationMethod,
)


class EdgeWiseModelConfig:
    """Edgewise Model Configuration class"""

    _java_class = 'oracle.pgx.config.mllib.EdgeWiseModelConfig'

    def __init__(self, java_edgewise_model_config) -> None:
        self._config = java_edgewise_model_config

    @property
    def shuffle(self) -> bool:
        """Whether or not shuffling is enabled."""
        return self._config.isShuffle()

    @property
    def input_feature_dim(self) -> int:
        """Get the input feature dimension."""
        return self._config.getInputFeatureDim()

    @property
    def edge_input_feature_dim(self) -> int:
        """Get the edge input feature dimension."""
        return self._config.getEdgeInputFeatureDim()

    @property
    def is_fitted(self) -> bool:
        """Whether or not the model is fitted."""
        return self._config.isFitted()

    @property
    def training_loss(self) -> float:
        """Get the training loss."""
        return self._config.getTrainingLoss()

    @property
    def batch_size(self) -> int:
        """Get the batch size."""
        return self._config.getBatchSize()

    @property
    def num_epochs(self) -> int:
        """Get the number of epochs."""
        return self._config.getNumEpochs()

    @property
    def learning_rate(self) -> float:
        """Get the learning rate."""
        return self._config.getLearningRate()

    @property
    def weight_decay(self) -> float:
        """Get the weight decay."""
        return self._config.getWeightDecay()

    @property
    def embedding_dim(self) -> int:
        """Get the embedding dimension."""
        return self._config.getEmbeddingDim()

    @property
    def seed(self) -> int:
        """Get the seed."""
        return self._config.getSeed()

    @property
    def conv_layer_configs(
        self,
    ) -> List[Union[GraphWiseConvLayerConfig, GraphWiseAttentionLayerConfig]]:
        """Get the conv layer configs."""
        return self.get_conv_layer_configs()

    @property
    def vertex_input_property_configs(self) -> Dict[str, InputPropertyConfig]:
        """Get the vertex input property configs."""
        java_vertex_input_property_configs = java_handler(
            self._config.getVertexInputPropertyConfigs, [])
        return _java_input_property_list_to_python_configs(
            java_vertex_input_property_configs
        )

    @property
    def edge_input_property_configs(self) -> Dict[str, InputPropertyConfig]:
        """Get the edge input property configs."""
        java_edge_input_property_configs = java_handler(
            self._config.getEdgeInputPropertyConfigs, [])
        return _java_input_property_list_to_python_configs(
            java_edge_input_property_configs
        )

    @property
    def vertex_input_property_names(self) -> Optional[List[str]]:
        """Get the vertex input property names."""
        names = self._config.getVertexInputPropertyNames()
        if not names:
            return None
        return names.toArray()

    @property
    def edge_input_property_names(self) -> Optional[List[str]]:
        """Get the edge input property names."""
        names = self._config.getEdgeInputPropertyNames()
        if not names:
            return None
        return names.toArray()

    @property
    def standardize(self) -> bool:
        """Whether or not standardization is enabled."""
        return self._config.isStandardize()

    @property
    def normalize(self) -> bool:
        """Whether or not normalization is enabled."""
        return self._config.isNormalize()

    @property
    def backend(self) -> str:
        """Get the backend."""
        return conversion.enum_to_python_str(self._config.getBackend())

    @property
    def edge_combination_method(self) -> EdgeCombinationMethod:
        """Get the edge combination method."""
        java_edge_combination_method = self._config.getEdgeCombinationMethod()
        params = {
            'use_source_vertex': java_edge_combination_method.isUseSourceVertex(),
            'use_destination_vertex': java_edge_combination_method.isUseDestinationVertex(),
            'use_edge': java_edge_combination_method.isUseEdge(),
        }
        method = java_edge_combination_method.getAggregationType().name()
        if method == "CONCATENATION":
            return ConcatEdgeCombinationMethod(**params)
        elif method == "PRODUCT":
            return ProductEdgeCombinationMethod(**params)
        else:
            raise ValueError

    def get_conv_layer_configs(
        self,
    ) -> List[Union[GraphWiseConvLayerConfig, GraphWiseAttentionLayerConfig]]:
        """Return a list of conv layer configs"""
        java_conv_layer_configs = java_handler(self._config.getConvLayerConfigs, [])
        return _java_conv_layer_configs_to_python_conv_layer_configs(java_conv_layer_configs)

    def set_batch_size(self, batch_size: int) -> None:
        """Set the batch size

        :param batch_size: batch size
        :type batch_size: int
        """
        java_handler(self._config.setBatchSize, [batch_size])

    def set_num_epochs(self, num_epochs: int) -> None:
        """Set the number of epochs

        :param num_epochs: number of epochs
        :type num_epochs: int
        """
        java_handler(self._config.setNumEpochs, [num_epochs])

    def set_learning_rate(self, learning_rate: float) -> None:
        """Set the learning rate

        :param learning_rate: initial learning rate
        :type learning_rate: int
        """
        java_handler(self._config.setLearningRate, [learning_rate])

    def set_weight_decay(self, weight_decay: float) -> None:
        """Set the weight decay

        :param weight_decay: weight decay
        :type weight_decay: float
        """
        java_handler(self._config.setWeightDecay, [weight_decay])

    def set_embedding_dim(self, embedding_dim: int) -> None:
        """Set the embedding dimension

        :param embedding_dim: embedding dimension
        :type embedding_dim: int
        """
        java_handler(self._config.setEmbeddingDim, [embedding_dim])

    def set_seed(self, seed: int) -> None:
        """Set the seed

        :param seed: seed
        :type seed: int
        """
        java_handler(self._config.setSeed, [seed])

    def set_fitted(self, fitted: bool) -> None:
        """Set the fitted flag

        :param fitted: fitted flag
        :type fitted: bool
        """
        java_handler(self._config.setFitted, [fitted])

    def set_shuffle(self, shuffle: bool) -> None:
        """Set the shuffling flag

        :param shuffle: shuffling flag
        :type shuffle: bool
        """
        java_handler(self._config.setShuffle, [shuffle])

    def set_training_loss(self, training_loss: float) -> None:
        """Set the training loss

        :param training_loss: training loss
        :type training_loss: float
        """
        java_handler(self._config.setTrainingLoss, [training_loss])

    def set_input_feature_dim(self, input_feature_dim: int) -> None:
        """Set the input feature dimension

        :param input_feature_dim: input feature dimension
        :type input_feature_dim: int
        """
        java_handler(self._config.setInputFeatureDim, [input_feature_dim])

    def set_edge_input_feature_dim(self, edge_input_feature_dim: int) -> None:
        """Set the edge input feature dimension

        :param edge_input_feature_dim: edge input feature dimension
        :type edge_input_feature_dim: int
        """
        java_handler(self._config.setEdgeInputFeatureDim, [edge_input_feature_dim])

    def set_standardize(self, standardize: bool) -> None:
        """Set the standardize flag

        :param standardize: standardize flag
        :type standardize: bool
        """
        java_handler(self._config.setStandardize, [standardize])

    def set_normalize(self, normalize: bool) -> None:
        """Whether or not normalization is enabled."""
        java_handler(self._config.setNormalize, [normalize])

    def set_edge_combination_method(self, edge_combination_method: EdgeCombinationMethod) -> None:
        """Set the edge combination method

        :param edge_combination_method: edge combination method
        :type edge_combination_method: EdgeCombinationMethod
        """
        java_handler(self._config.setEdgeCombinationMethod, [edge_combination_method._java_obj])
