# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_basic_plots.ipynb.

# %% auto 0
__all__ = ['data_histplot', 'data_ecdfplot', 'hist_ecdf_plots', 'plot_stats', 'plot_hist_stats', 'plot_data_collection',
           'show_fundus']

# %% ../nbs/01_basic_plots.ipynb 3
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from typing import Optional
from PIL import Image
from smart_open import open

# %% ../nbs/01_basic_plots.ipynb 4
from pheno_utils.config import (
    ALL_COLOR,
    FEMALE_COLOR, 
    MALE_COLOR, 
    generate_synthetic_data
    )

# %% ../nbs/01_basic_plots.ipynb 5
def data_histplot(data: pd.DataFrame, col: str, feature_str: Optional[str] = None,
                  gender_col: str = "sex", plot_both_genders: bool=True, ax=None) -> None:
    """
    Plots a histogram from a DataFrame for a specific column.

    Args:
        data (pd.DataFrame): The DataFrame containing the data to plot.
        col (str): The name of the column to plot.
        feature_str (Optional[str], optional): The name of the feature to plot. Defaults to None.
        gender_col (str, optional): The name of the column containing gender information. Defaults to "sex".
        plot_both_genders (bool, optional): Whether to plot both genders or just one. Defaults to True.
        ax ([type], optional): The axis to plot on. Defaults to None.
    """
    try:
        assert col in data.columns, f"Column {col} not found in data"
        if plot_both_genders:
            assert gender_col in data.columns, f"Gender column {gender_col} not found in data"
    except AssertionError as e:
        print(e)
        return

    if feature_str is None:
        feature_str = col

    if ax is None:
        _, ax = plt.subplots(1, 1, figsize=(6, 4))

    # Plot histogram(s)
    if plot_both_genders:
        idx_male = data[gender_col] == 1
        idx_female = data[gender_col] == 0
        sns.histplot(data.loc[idx_female, col], ax=ax, color=FEMALE_COLOR,
                    label=f'females (N={(data.loc[idx_female, col].count()):,})', alpha=0.5)
        sns.histplot(data.loc[idx_male, col], ax=ax, color=MALE_COLOR,
                    label=f'males (N={(data.loc[idx_male, col].count()):,})', alpha=0.5)
        
    else:
        sns.histplot(data.loc[:, col], ax=ax, color=ALL_COLOR,
                    label=f'N={(data.loc[:, col].count()):,}', alpha=0.5)

    ax.set_xlabel(feature_str, fontsize=16)
    ax.spines["right"].set_visible(False)
    ax.spines["top"].set_visible(False)
    ax.legend()

# %% ../nbs/01_basic_plots.ipynb 9
def data_ecdfplot(data: pd.DataFrame, col: str, feature_str: Optional[str] = None,
                  gender_col: str = "sex", plot_both_genders: bool=True, ax=None) -> None:
    """
    Plots an empirical cumulative distribution function (ECDF) from a DataFrame for a specific column.

    Args:
        data (pd.DataFrame): The DataFrame containing the data to plot.
        col (str): The name of the column to plot.
        feature_str (Optional[str], optional): The name of the feature to plot. Defaults to None.
        gender_col (str, optional): The name of the column containing gender information. Defaults to "sex".
        plot_both_genders (bool, optional): Whether to plot both genders or just one. Defaults to True.
        ax ([type], optional): The axis to plot on. Defaults to None.
    """
    try:
        assert col in data.columns, f"Column {col} not found in data"
        if plot_both_genders:
            assert gender_col in data.columns, f"Gender column {gender_col} not found in data"
    except AssertionError as e:
        print(e)
        return

    if feature_str is None:
        feature_str = col

    if ax is None:
        _, ax = plt.subplots(1, 1, figsize=(6, 4))

    # Plot ecdf(s)
    if plot_both_genders:
        idx_male = data[gender_col] == 1
        idx_female = data[gender_col] == 0
        sns.ecdfplot(data.loc[idx_female, col], ax=ax, color=FEMALE_COLOR,
                    label=f'females (N={(data.loc[idx_female, col].count()):,})', alpha=0.5, lw=4)
        sns.ecdfplot(data.loc[idx_male, col], ax=ax, color=MALE_COLOR,
                    label=f'males (N={(data.loc[idx_male, col].count()):,})', alpha=0.5, lw=4)
        
    else:
        sns.ecdfplot(data.loc[:, col], ax=ax, color=ALL_COLOR,
                    label=f'N={(data.loc[:, col].count()):,}', alpha=0.5, lw=4)

    ax.set_xlabel(feature_str, fontsize=16)
    ax.spines["right"].set_visible(False)
    ax.spines["top"].set_visible(False)
    ax.legend()

# %% ../nbs/01_basic_plots.ipynb 12
def hist_ecdf_plots(data: pd.DataFrame, col: str, feature_str: Optional[str] = None,
                  gender_col: str = "sex", plot_both_genders: bool=True) -> None:
    """
    Plots histograms and empirical cumulative distribution functions (ECDFs) from a DataFrame
    for a specific column.

    Args:
        data: The input DataFrame containing the data to plot.
        col: The column name to plot.
        feature_str: The title of the plot. If not provided, the column name will be used.
        gender_col: The column name indicating sex (default is 'sex' - female:0; male:1).
        plot_both_genders (bool, optional): Whether to plot both genders or just one. Defaults to True.

    Returns:
        None
    """
    try:
        assert col in data.columns, f"Column {col} not found in data"
        if plot_both_genders:
            assert gender_col in data.columns, f"Gender column {gender_col} not found in data"
    except AssertionError as e:
        print(e)
        return

    if feature_str is None:
        feature_str = col

    fig, axes = plt.subplots(1, 2, figsize=(10, 4))

    # Plot histograms
    ax = axes[0]
    data_histplot(data=data, col=col, feature_str=feature_str,
                  gender_col=gender_col, plot_both_genders=plot_both_genders, ax=ax)
    # Plot ECDFs
    ax = axes[1]
    data_ecdfplot(data=data, col=col, feature_str=feature_str,
                  gender_col=gender_col, plot_both_genders=plot_both_genders, ax=ax)

    fig.tight_layout()
    plt.show()

# %% ../nbs/01_basic_plots.ipynb 14
def plot_stats(data: pd.DataFrame, col: str, ax: plt.Axes, color: str,
               x_position: float = 0.3, label: Optional[str] = "All") -> None:
    """
    Adds a text box to an axis object with summary statistics for a given column in a pandas DataFrame.

    Args:
        data (pd.DataFrame): The pandas DataFrame containing the data to calculate statistics for.
        col (str): The name of the column to calculate statistics for.
        ax (plt.Axes): The axis object to add the text box to.
        color (str): The color of the text box.
        x_position (float, optional): The x position of the text box. Defaults to 0.3.
        label (Optional[str], optional): The label to display in the text box. Defaults to "All".
    """
    # Calculate statistics
    stats = data[col].describe(percentiles=[.10, .90])

    # Rename "50%" to "Median" and "count" to "N"
    stats = stats.rename(index={"50%": "Median", "count": "N"})

    # Ensure "N" is an integer value
    stats["N"] = int(stats["N"])

    # Format statistics for display
    stats_str = f"{label}\n--------------------\n"
    for k, v in stats.items():
        if k == "N":
            stats_str += f"{k}: {int(v)}\n\n"  # Handle "N" separately
        elif np.issubdtype(type(v), np.integer):
            stats_str += f"{k}: {v}\n\n"
        else:
            stats_str += f"{k}: {round(v, 2)}\n\n"

    # Remove trailing newline
    stats_str = stats_str.rstrip("\n")

    # Add textbox to axes
    ax.text(x_position, 0.5, stats_str,
            horizontalalignment='center',
            verticalalignment='center',
            fontsize=9, color='black',
            bbox={'facecolor': color, 'edgecolor': 'black', 'alpha': 0.5, 'pad': 8},
            transform=ax.transAxes)

    # Hide axes
    ax.axis('off')

# %% ../nbs/01_basic_plots.ipynb 15
def plot_hist_stats(data: pd.DataFrame, col: str, feature_str: Optional[str] = None,
                    gender_col: str = "sex", plot_both_genders: bool=True) -> None:
    """
    Plots a histogram of a given column in a pandas DataFrame and adds summary statistics to the plot.

    Args:
        data (pd.DataFrame): The pandas DataFrame containing the data to plot.
        col (str): The name of the column to plot.
        feature_str (Optional[str], optional): A string describing the feature being plotted. Defaults to None.
        gender_col (str, optional): The name of the column containing gender information. Defaults to "sex".
        plot_both_genders (bool, optional): Whether to plot statistics separately for males and females. Defaults to True.
    """
    # Create a figure and an axes
    fig, axes = plt.subplots(1,2, figsize=(10, 4))

    # Plot histogram on first axis
    ax = axes[0]
    data_histplot(data, col, feature_str, gender_col, plot_both_genders, ax=ax)
    ax.get_legend().remove()  # Remove legend from histogram

    # Plot statistics on second axis
    ax = axes[1]
    if plot_both_genders:
        # Separate data by gender
        idx_male = data[gender_col] == 1
        idx_female = data[gender_col] == 0

        # Plot statistics for males and females separately
        plot_stats(data.loc[idx_male], col, ax=ax, color=MALE_COLOR, x_position=0, label="Males")
        plot_stats(data.loc[idx_female], col, ax=ax, color=FEMALE_COLOR, x_position=0.4, label="Females")

        # Plot statistics for all data
        plot_stats(data.loc[:], col, ax=ax, color=ALL_COLOR, x_position=0.8)

    else:
        # Plot statistics for all data
        plot_stats(data, col, ax=ax, color=ALL_COLOR, x_position=0)

# %% ../nbs/01_basic_plots.ipynb 19
def plot_data_collection(data: pd.DataFrame, date_col: str = "collection_date", feature_str: Optional[str] = None, ax: Optional[plt.Axes] = None) -> None:
    """
    Plots a histogram of the specified column in a pandas DataFrame and excludes the last point from the plot.

    Args:
        data (pd.DataFrame): The pandas DataFrame containing the data to plot.
        date_col (str, optional): The name of the column containing the dates. Defaults to "collection_date".
        feature_str (Optional[str], optional): The name of the feature to plot. If None, the name of the date column will be used. Defaults to None.
        ax (Optional[plt.Axes], optional): The axis object to plot on. If None, a new figure and axis will be created. Defaults to None.
    """
    if feature_str is None:
        feature_str = ""

    if ax is None:
        _, ax = plt.subplots(1, 1, figsize=(6, 3))

    df = data.sort_values(date_col).dropna()

    # Calculate cumulative count of rows
    df['cumulative_count'] = range(1, len(df) + 1)

    # Plot cumulative count over time
    ax.plot(df[date_col], df['cumulative_count'], lw=3, alpha=0.8)

    # Rotate x-axis tick labels by 90 degrees
    plt.xticks(rotation=90)

    # Remove top and right spines
    ax.spines["right"].set_visible(False)
    ax.spines["top"].set_visible(False)

    ax.set_ylabel(f"{feature_str} data collected", fontsize=14)

# %% ../nbs/01_basic_plots.ipynb 21
def show_fundus(fname: str) -> None:
    """
    Display a fundus image from an input file path.
    Args:
        fname (str): The file path to the fundus image.
    """
    fig, ax = plt.subplots(1, 1, figsize=(6, 6))
    img = Image.open(open(fname, 'rb'))
    ax.imshow(img, cmap="gray")
    ax.set_xticks([])
    ax.set_yticks([])
    ax.axis('off')
