# Copyright 2026 Ali Seyhan, Baki Turhan
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# 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.

"""
QuinkGL: Decentralized Gossip Learning Framework

A Python framework for decentralized machine learning using
gossip-based peer-to-peer communication.

Example:
    from quinkgl import (
        GossipNode,
        PyTorchModel,
        RandomTopology,
        FedAvg,
        TrainingConfig
    )

    model = PyTorchModel(my_pytorch_model)
    node = GossipNode(
        node_id="alice",
        domain="health",
        model=model,
        topology=RandomTopology(),
        aggregation=FedAvg()
    )
    await node.start()
    await node.run_continuous(training_data)
"""

__version__ = "0.1.0"

# =============================================================================
# CORE - Main node classes
# =============================================================================
from quinkgl.core.learning_node import LearningNode, GLNode
from quinkgl.network.gossip_node import GossipNode

# =============================================================================
# MODELS - Framework-specific model wrappers
# =============================================================================
from quinkgl.models.base import (
    ModelWrapper,
    TrainingConfig,
    TrainingResult
)
from quinkgl.models.pytorch import PyTorchModel

# TensorFlow is optional - only import if available
try:
    from quinkgl.models.tensorflow import TensorFlowModel
    _tensorflow_available = True
except ImportError:
    _tensorflow_available = False
    TensorFlowModel = None  # type: ignore

# =============================================================================
# TOPOLOGY - Peer selection strategies
# =============================================================================
from quinkgl.topology.base import (
    TopologyStrategy,
    PeerInfo,
    SelectionContext
)
from quinkgl.topology.random import RandomTopology
from quinkgl.topology.cyclon import CyclonTopology

# =============================================================================
# AGGREGATION - Model combining strategies
# =============================================================================
from quinkgl.aggregation.base import (
    AggregationStrategy,
    ModelUpdate,
    AggregatedModel
)
from quinkgl.aggregation.fedavg import FedAvg
from quinkgl.aggregation.strategies import (
    FedProx,
    FedAvgM,
    TrimmedMean,
    Krum,
    MultiKrum
)

# =============================================================================
# GOSSIP - Model aggregation orchestration
# =============================================================================
from quinkgl.gossip.aggregator import ModelAggregator

# =============================================================================
# DATA - Dataset loading and splitting (OPTIONAL - not included in pip package)
# =============================================================================
# Data loading utilities are in scripts/run_gossip_node.py for full functionality
# or use torchvision directly for CIFAR-10
try:
    from quinkgl.data import (
        DatasetLoader,
        FederatedDataSplitter,
        DatasetInfo
    )
    _data_available = True
except (ImportError, ModuleNotFoundError):
    # Data module not available in pip package
    DatasetLoader = None  # type: ignore
    FederatedDataSplitter = None  # type: ignore
    DatasetInfo = None  # type: ignore
    _data_available = False

# =============================================================================
# MCP - Model Context Protocol for monitoring (optional)
# =============================================================================
try:
    from quinkgl.mcp import MetricsCollector, MCPServer
    _mcp_available = True
except ImportError:
    _mcp_available = False
    MetricsCollector = None  # type: ignore
    MCPServer = None  # type: ignore

# =============================================================================
# PUBLIC API
# =============================================================================
__all__ = [
    # Core
    "LearningNode",
    "GLNode",  # Backward compatibility alias
    "GossipNode",
    
    # Models
    "ModelWrapper",
    "TrainingConfig",
    "TrainingResult",
    "PyTorchModel",
    "TensorFlowModel",
    
    # Topology
    "TopologyStrategy",
    "RandomTopology",
    "CyclonTopology",
    "PeerInfo",
    "SelectionContext",
    
    # Aggregation
    "AggregationStrategy",
    "FedAvg",
    "FedProx",
    "FedAvgM",
    "TrimmedMean",
    "Krum",
    "MultiKrum",
    "ModelUpdate",
    "AggregatedModel",
    
    # Gossip
    "ModelAggregator",
    
    # Data
    "DatasetLoader",
    "FederatedDataSplitter",
    "DatasetInfo",
    
    # MCP (optional)
    "MetricsCollector",
    "MCPServer",
]

