from __future__ import annotations

import logging
from typing import Any, Generator

import pandas as pd
from pydantic import BaseModel, Field, PositiveFloat, PositiveInt

from graph.connections.db_connection import DbConnection
from graph.exploration.algorithms.backend_service import (
    GraphAnalyticsExecutionLocation,
    GraphAnalyticsSamplingMode,
    get_graph_analytics_backend_service,
)
from graph.exploration.algorithms.projected_graph import ProjectedGraph

logger = logging.getLogger(__name__)

class PageRankParams(BaseModel):
    damping_factor: float = Field(..., ge=0.0, le=1.0)
    max_iterations: PositiveInt
    tolerance: PositiveFloat
    normalizeInitial: bool


def compute_page_rank(
    conn: DbConnection,
    projected_graph: ProjectedGraph,
    params: PageRankParams,
    batch_size: int = 10000,
    *,
    location: GraphAnalyticsExecutionLocation = GraphAnalyticsExecutionLocation.NATIVE,
    sampling_mode: GraphAnalyticsSamplingMode | None = None,
    sampling_limit: int | None = None,
) -> Generator[pd.DataFrame, Any, None]:
    return get_graph_analytics_backend_service().execute(
        "page_rank",
        conn,
        projected_graph,
        params,
        batch_size=batch_size,
        location=location,
        sampling_mode=sampling_mode,
        sampling_limit=sampling_limit,
    )
