from dataclasses import dataclass
from typing import Literal, Tuple, Union, Optional, Dict, Any
import numpy as np
import networkx as nx

NetworkFileResult = Tuple[str, Dict[Any, int], bool] # (file_path, node_map, redundant)

NetworkFormatEdgelist = Tuple[
    Literal["edgelist"], str, Optional[bool], Optional[str], Optional[str], Optional[bool]
]
# ("edgelist", path, redundant, delimiter, comment, cache)

NetworkFormatFortranEdgelist = Tuple[
    Literal["fortran-edgelist"], str, Optional[bool]
]
# ("fortran-edgelist", path, redundant)

# TODO: get URLs from OneDrive
#NetworkFormatPL = Tuple[
#    Literal["PL"],
#    Literal[2.5, 3.0, 6.0],
#    Literal[100, 1000, 10000, 100000, 1000000],
#    Optional[Literal[1, 2, 3, 4, 5]]
#]
## ("PL", gamma, N, sample)

NetworkFormatNetworkX = Tuple[
    Literal["networkx"], Union[nx.Graph, nx.DiGraph, nx.MultiGraph, nx.MultiDiGraph], Optional[bool]
]
# ("networkx", graph, cache)

NetworkFormat = Union[
    NetworkFormatEdgelist, # Nodes can be anything
    NetworkFormatFortranEdgelist, # Nodes are indexed from 1 to N
    #NetworkFormatPL, # Will download from datasets
    NetworkFormatNetworkX # Generated with NetworkX
]

@dataclass
class TemporalResult:
    t: np.ndarray
    rho_avg: np.ndarray
    rho_var: np.ndarray
    n_samples: int

@dataclass
class SimulationResult:
    network: NetworkFormat
    node_map: dict
    temporal: TemporalResult
    active_states: Optional[dict] = None  # {sample_id: {time: {"nodes": [...]}}}
    nx_graph: Union[nx.Graph, nx.DiGraph, nx.MultiGraph, nx.MultiDiGraph, None] = None
    stdout: Optional[str] = None
    stderr: Optional[str] = None

@dataclass
class SimulationArgs:
    """
    Simulation arguments for dynSIS

    Attributes
    ----------
    verbose : bool
        Enable verbose output.
    verbose_level : str
        Logging level: 'info', 'warning', 'error', 'debug'.
    seed : int
        Random seed.
    remove_files : bool
        Remove temporary files after execution.
    network: NetworkFormat
        Network specification as a tuple ([optional parameters in brackets]):
        ("edgelist", path, [redundant], [delimiter], [comment], [cache])
        ("fortran-edgelist", path, [redundant])
        Not implemented yet: ("PL", gamma, N, [sample])
        ("networkx", graph, [cache])
    output_dir : Optional[str]
        Directory to store simulation output. If None, uses temporary folder.
    algorithm : str
        Simulation algorithm: 'SIS-OGA'.
    sampler : str
        Sampling method: 'rejection_maxheap' or 'btree'.
    tmax : int
        Maximum simulation time.
    use_qs : bool
        Whether to use quasi-stationary method.
    n_samples : int
        Number of samples per simulation.
    time_scale : str
        Temporal scale for output: 'uniform' or 'powerlaw'.
    initial_condition : tuple
        Initial state specification:
        ('fraction', float) -> fraction of infected nodes
        ('number', int) -> exact number of initially infected nodes
    export_states : bool
        Whether to export the full state trajectory.
    build_nx_graph : bool
        Whether to build and return the NetworkX graph representation of the network.
    """
    # General
    verbose: bool = True
    verbose_level: Literal["info", "warning", "error", "debug"] = "warning"
    seed: int = 42
    remove_files: bool = False

    # IO
    # network: NetworkFormat = ("PL", # default
    #                           3.0,
    #                           100,
    #                           1
    #                         )
    network: NetworkFormat = ("networkx", # default
                              nx.barabasi_albert_graph(100, 3),
                              False
                            )
    output_dir: Optional[str] = None
    # If True, existing output files in `output_dir` (results/dat and states_nodes)
    # will be removed before running the simulation. If False (default), the
    # runner will raise an error when these files are detected to avoid
    # accidental overwrites.
    overwrite_output: bool = False

    # Algorithm
    algorithm: Literal["SIS-OGA"] = "SIS-OGA"
    sampler: Literal["rejection_maxheap", "btree"] = "btree"

    # Dynamics
    tmax: int = 100
    use_qs: bool = False
    n_samples: int = 10
    time_scale: Literal["uniform", "powerlaw"] = "uniform"
    initial_condition: Union[Tuple[Literal["fraction"], float], Tuple[Literal["number"], int]] = ("fraction", 1.0)

    # IO export
    export_states: bool = False
    build_nx_graph: bool = False
