"""Documentation for commonly used functions in pandas, sklearn, and pgmpy."""

from __future__ import annotations

from typing import Dict, List


def get_pandas_docs() -> Dict[str, str]:
    """Return documentation for pandas functions."""
    return {
        "pd.DataFrame": """
pd.DataFrame(data, index=None, columns=None)
Create a two-dimensional, size-mutable, potentially heterogeneous tabular data.

Parameters:
- data: dict, list, or DataFrame
- index: Index or array-like (optional)
- columns: Index or array-like (optional)

Returns: DataFrame object
""",
        "df.head()": """
df.head(n=5)
Return the first n rows.

Parameters:
- n: int, default 5

Returns: DataFrame
""",
        "df.tail()": """
df.tail(n=5)
Return the last n rows.

Parameters:
- n: int, default 5

Returns: DataFrame
""",
        "df.shape": """
df.shape
Return a tuple representing the dimensionality of the DataFrame.

Returns: tuple of (rows, columns)
""",
        "df.describe()": """
df.describe(percentiles=None, include=None, exclude=None)
Generate descriptive statistics.

Returns: DataFrame with summary statistics
""",
        "df.loc[]": """
df.loc[row_indexer, column_indexer]
Access a group of rows and columns by label(s) or a boolean array.

Parameters:
- row_indexer: label, array of labels, or slice
- column_indexer: label, array of labels, or slice

Returns: Series or DataFrame
""",
        "df.iloc[]": """
df.iloc[row_indexer, column_indexer]
Purely integer-location based indexing for selection by position.

Parameters:
- row_indexer: int, list of ints, or slice
- column_indexer: int, list of ints, or slice

Returns: Series or DataFrame
""",
        "df.drop()": """
df.drop(labels=None, axis=0, index=None, columns=None, inplace=False)
Drop specified labels from rows or columns.

Parameters:
- labels: single label or list-like
- axis: {0 or 'index', 1 or 'columns'}, default 0
- inplace: bool, default False

Returns: DataFrame or None
""",
        "df.dropna()": """
df.dropna(axis=0, how='any', thresh=None, subset=None, inplace=False)
Remove missing values.

Parameters:
- axis: {0 or 'index', 1 or 'columns'}, default 0
- how: {'any', 'all'}, default 'any'
- inplace: bool, default False

Returns: DataFrame or None
""",
        "pd.read_csv()": """
pd.read_csv(filepath_or_buffer, sep=',', header='infer', names=None, index_col=None)
Read a comma-separated values (csv) file into DataFrame.

Parameters:
- filepath_or_buffer: str, path object or file-like object
- sep: str, default ','
- header: int, list of int, default 'infer'
- names: array-like, optional

Returns: DataFrame
""",
        "df.to_string()": """
df.to_string(buf=None, columns=None, col_space=None, header=True, index=True)
Render a DataFrame to a console-friendly tabular output.

Returns: str
""",
    }


def get_sklearn_docs() -> Dict[str, str]:
    """Return documentation for scikit-learn functions."""
    return {
        "train_test_split()": """
train_test_split(*arrays, test_size=None, train_size=None, random_state=None, shuffle=True, stratify=None)
Split arrays or matrices into random train and test subsets.

Parameters:
- *arrays: sequence of indexables with same length
- test_size: float or int, default None
- train_size: float or int, default None
- random_state: int, RandomState instance or None
- shuffle: bool, default True
- stratify: array-like or None, default None

Returns: list of split arrays
""",
        "LabelEncoder()": """
LabelEncoder()
Encode target labels with value between 0 and n_classes-1.

Methods:
- fit(y): Fit label encoder
- transform(y): Transform labels to normalized encoding
- fit_transform(y): Fit label encoder and return encoded labels
- inverse_transform(y): Transform labels back to original encoding
""",
        "StandardScaler()": """
StandardScaler(copy=True, with_mean=True, with_std=True)
Standardize features by removing the mean and scaling to unit variance.

Methods:
- fit(X): Compute the mean and std to be used for later scaling
- transform(X): Perform standardization by centering and scaling
- fit_transform(X): Fit to data, then transform it
- inverse_transform(X): Scale back the data to the original representation
""",
        "DecisionTreeClassifier()": """
DecisionTreeClassifier(criterion='gini', splitter='best', max_depth=None, min_samples_split=2, min_samples_leaf=1, random_state=None)
A decision tree classifier.

Parameters:
- criterion: {'gini', 'entropy', 'log_loss'}, default 'gini'
- max_depth: int, default None
- min_samples_split: int or float, default 2
- random_state: int, RandomState instance or None

Methods:
- fit(X, y): Build a decision tree classifier from the training set
- predict(X): Predict class or regression value for X
- score(X, y): Return the mean accuracy on the given test data
""",
        "KNeighborsClassifier()": """
KNeighborsClassifier(n_neighbors=5, weights='uniform', algorithm='auto', leaf_size=30)
Classifier implementing the k-nearest neighbors vote.

Parameters:
- n_neighbors: int, default 5
- weights: {'uniform', 'distance'}, default 'uniform'
- algorithm: {'auto', 'ball_tree', 'kd_tree', 'brute'}, default 'auto'

Methods:
- fit(X, y): Fit the k-nearest neighbors classifier
- predict(X): Predict the class labels for the provided data
- predict_proba(X): Return probability estimates for the test data
""",
        "SVC()": """
SVC(C=1.0, kernel='rbf', degree=3, gamma='scale', random_state=None)
C-Support Vector Classification.

Parameters:
- C: float, default 1.0 (regularization parameter)
- kernel: {'linear', 'poly', 'rbf', 'sigmoid'}, default 'rbf'
- gamma: {'scale', 'auto'} or float, default 'scale'
- random_state: int, RandomState instance or None

Methods:
- fit(X, y): Fit the SVM model
- predict(X): Perform classification on samples in X
- score(X, y): Return the mean accuracy on the given test data
""",
        "RandomForestClassifier()": """
RandomForestClassifier(n_estimators=100, criterion='gini', max_depth=None, min_samples_split=2, random_state=None)
A random forest classifier.

Parameters:
- n_estimators: int, default 100
- criterion: {'gini', 'entropy', 'log_loss'}, default 'gini'
- max_depth: int, default None
- random_state: int, RandomState instance or None

Methods:
- fit(X, y): Build a forest of trees from the training set
- predict(X): Predict class for X
- predict_proba(X): Predict class probabilities for X
""",
        "GaussianNB()": """
GaussianNB(priors=None, var_smoothing=1e-09)
Gaussian Naive Bayes (GaussianNB).

Parameters:
- priors: array-like of shape (n_classes,), default None
- var_smoothing: float, default 1e-09

Methods:
- fit(X, y): Fit Gaussian Naive Bayes according to X, y
- predict(X): Perform classification on an array of test vectors X
- predict_proba(X): Return probability estimates for the test vector X
""",
        "BaggingClassifier()": """
BaggingClassifier(estimator=None, n_estimators=10, max_samples=1.0, bootstrap=True, oob_score=False, random_state=None)
A Bagging classifier.

Parameters:
- estimator: object, default None
- n_estimators: int, default 10
- max_samples: int or float, default 1.0
- bootstrap: bool, default True
- oob_score: bool, default False
- random_state: int, RandomState instance or None

Methods:
- fit(X, y): Build a Bagging ensemble of estimators from the training set
- predict(X): Predict class for X
""",
        "AdaBoostClassifier()": """
AdaBoostClassifier(estimator=None, n_estimators=50, learning_rate=1.0, algorithm='SAMME.R', random_state=None)
An AdaBoost classifier.

Parameters:
- estimator: object, default None
- n_estimators: int, default 50
- learning_rate: float, default 1.0
- algorithm: {'SAMME', 'SAMME.R'}, default 'SAMME.R'
- random_state: int, RandomState instance or None

Methods:
- fit(X, y): Build a boosted classifier from the training set
- predict(X): Predict classes for X
- predict_proba(X): Predict class probabilities for X
""",
        "PCA()": """
PCA(n_components=None, copy=True, whiten=False, svd_solver='auto', random_state=None)
Principal component analysis (PCA).

Parameters:
- n_components: int, float or 'mle', default None
- copy: bool, default True
- random_state: int, RandomState instance or None

Methods:
- fit(X): Fit the model with X
- transform(X): Apply dimensionality reduction to X
- fit_transform(X): Fit the model with X and apply the dimensionality reduction
""",
        "GaussianMixture()": """
GaussianMixture(n_components=1, covariance_type='full', random_state=None)
Gaussian Mixture Model (GMM).

Parameters:
- n_components: int, default 1
- covariance_type: {'full', 'tied', 'diag', 'spherical'}, default 'full'
- random_state: int, RandomState instance or None

Methods:
- fit(X): Estimate model parameters with the EM algorithm
- predict(X): Predict the labels for the data samples in X
- score(X): Compute the per-sample average log-likelihood
""",
        "KMeans()": """
KMeans(n_clusters=8, init='k-means++', n_init=10, max_iter=300, random_state=None)
K-Means clustering.

Parameters:
- n_clusters: int, default 8
- init: {'k-means++', 'random'}, default 'k-means++'
- n_init: int, default 10
- random_state: int, RandomState instance or None

Methods:
- fit(X): Compute k-means clustering
- fit_predict(X): Compute cluster centers and predict cluster index for each sample
- predict(X): Predict the closest cluster each sample in X belongs to
""",
        "accuracy_score()": """
accuracy_score(y_true, y_pred, normalize=True, sample_weight=None)
Accuracy classification score.

Parameters:
- y_true: 1d array-like, or label indicator array
- y_pred: 1d array-like, or label indicator array
- normalize: bool, default True

Returns: float (if normalize=True) or int (if normalize=False)
""",
        "confusion_matrix()": """
confusion_matrix(y_true, y_pred, labels=None, sample_weight=None)
Compute confusion matrix to evaluate the accuracy of a classification.

Parameters:
- y_true: array-like of shape (n_samples,)
- y_pred: array-like of shape (n_samples,)
- labels: array-like of shape (n_classes,), default None

Returns: ndarray of shape (n_classes, n_classes)
""",
        "classification_report()": """
classification_report(y_true, y_pred, target_names=None, output_dict=False)
Build a text report showing the main classification metrics.

Parameters:
- y_true: 1d array-like, or label indicator array
- y_pred: 1d array-like, or label indicator array
- target_names: list of strings, default None
- output_dict: bool, default False

Returns: str or dict
""",
        "cross_val_score()": """
cross_val_score(estimator, X, y=None, scoring=None, cv=None, n_jobs=None)
Evaluate a score by cross-validation.

Parameters:
- estimator: estimator object
- X: array-like of shape (n_samples, n_features)
- y: array-like of shape (n_samples,), default None
- cv: int, cross-validation generator or an iterable, default None
- scoring: str, callable, list/tuple or dict, default None

Returns: array of scores
""",
        "roc_curve()": """
roc_curve(y_true, y_score, pos_label=None, sample_weight=None, drop_intermediate=True)
Compute Receiver operating characteristic (ROC) curve.

Parameters:
- y_true: ndarray of shape (n_samples,)
- y_score: ndarray of shape (n_samples,)
- pos_label: int or str, default None

Returns: fpr, tpr, thresholds (arrays)
""",
        "roc_auc_score()": """
roc_auc_score(y_true, y_score, average='macro', sample_weight=None)
Compute Area Under the Receiver Operating Characteristic Curve (ROC AUC).

Parameters:
- y_true: array-like of shape (n_samples,)
- y_score: array-like of shape (n_samples,)
- average: {'micro', 'macro', 'samples', 'weighted'}, default 'macro'

Returns: float
""",
        "ConfusionMatrixDisplay()": """
ConfusionMatrixDisplay(confusion_matrix, display_labels=None)
Confusion Matrix visualization.

Parameters:
- confusion_matrix: ndarray of shape (n_classes, n_classes)
- display_labels: array-like of shape (n_classes,), default None

Methods:
- plot(cmap='viridis', xticks_rotation='horizontal'): Plot visualization
""",
        "export_graphviz()": """
export_graphviz(decision_tree, out_file=None, feature_names=None, class_names=None, filled=False, rounded=True)
Export a decision tree in DOT format.

Parameters:
- decision_tree: decision tree classifier
- out_file: object or str, default None
- feature_names: list of strings, default None
- class_names: list of strings, default None
- filled: bool, default False
- rounded: bool, default True

Returns: str
""",
    }


def get_pgmpy_docs() -> Dict[str, str]:
    """Return documentation for pgmpy functions."""
    return {
        "DiscreteBayesianNetwork()": """
DiscreteBayesianNetwork(ebunch=None, latents=set())
Base class for Bayesian Network.

Parameters:
- ebunch: list of edges, default None
- latents: set, default set()

Methods:
- add_cpds(*cpds): Add CPDs (Conditional Probability Distributions) to the model
- check_model(): Check the model for various errors
- get_cpds(): Returns the CPDs that have been added till now
""",
        "TabularCPD()": """
TabularCPD(variable, variable_card, values, evidence=None, evidence_card=None)
Represents Tabular Conditional Probability Distribution.

Parameters:
- variable: int or string (name of the variable)
- variable_card: int (cardinality of variable)
- values: array-like, shape (variable_card, 1) or (variable_card, evidence_card)
- evidence: list, default None (list of variable names which are evidences)
- evidence_card: list, default None (list of cardinalities of evidence variables)

Returns: TabularCPD object
""",
        "VariableElimination()": """
VariableElimination(model)
Variable Elimination for inference in Bayesian Networks.

Parameters:
- model: BayesianNetwork instance

Methods:
- query(variables, evidence=None, elimination_order=None): Query using Variable Elimination
  - variables: list (list of variables for which you want to compute the probability)
  - evidence: dict, default None (dict of variable: value pairs)
  
Returns: Factor object
""",
    }


def print_all_docs() -> None:
    """Print all documentation for pandas, sklearn, and pgmpy."""
    sections = [
        ("=" * 80, "PANDAS DOCUMENTATION", get_pandas_docs()),
        ("=" * 80, "SCIKIT-LEARN (SKLEARN) DOCUMENTATION", get_sklearn_docs()),
        ("=" * 80, "PGMPY DOCUMENTATION", get_pgmpy_docs()),
    ]
    
    for separator, title, docs_dict in sections:
        print(separator)
        print(title)
        print(separator)
        print()
        
        for func_name, doc_string in docs_dict.items():
            print(f"### {func_name}")
            print(doc_string.strip())
            print()
        
        print()


def get_docs_by_library(library: str) -> Dict[str, str]:
    """Get documentation for a specific library.
    
    Args:
        library: One of 'pandas', 'sklearn', or 'pgmpy'
        
    Returns:
        Dictionary mapping function names to their documentation
    """
    library = library.lower()
    if library == "pandas":
        return get_pandas_docs()
    elif library in ("sklearn", "scikit-learn"):
        return get_sklearn_docs()
    elif library == "pgmpy":
        return get_pgmpy_docs()
    else:
        raise ValueError(f"Unknown library: {library}. Choose from 'pandas', 'sklearn', or 'pgmpy'")


def print_docs_by_library(library: str) -> None:
    """Print documentation for a specific library.
    
    Args:
        library: One of 'pandas', 'sklearn', or 'pgmpy'
    """
    docs = get_docs_by_library(library)
    library_name = library.upper() if library.lower() == "sklearn" else library.upper()
    
    print("=" * 80)
    print(f"{library_name} DOCUMENTATION")
    print("=" * 80)
    print()
    
    for func_name, doc_string in docs.items():
        print(f"### {func_name}")
        print(doc_string.strip())
        print()

