import polars as pl
import numpy as np

def calc_scores_from_roc(
    df_roc: pl.DataFrame, n_neg: int, n_pos: int
) -> pl.DataFrame:
    """Derive observation-level scores and targets from points on the ROC curve

    Args:
        df_roc (pl.DataFrame): ROC curve characterized as dataset with (`fpr`,`tpr`) pairs
        n_neg (int): Global number of true negatives in target dataset
        n_pos (int): Global number of true positives in target dataset

    Returns:
        pl.DataFrame: Dataset containing `score` and `target` for each sample observation
    """

    df_scorebins = _gen_roc_to_scorebins(df_roc, n_neg, n_pos)
    df_scores = _gen_scorebins_to_scores(df_scorebins)
    return df_scores


def _gen_roc_to_scorebins(df_roc: pl.DataFrame, n_neg: int, n_pos: int) -> pl.DataFrame:
    """Derive score bins and frequencies of true positives and true negatives

    Args:
        df_roc (pl.DataFrame): ROC curve characterized as dataset with (`fpr`,`tpr`) pairs
        n_neg (int): Global number of true negatives in target dataset
        n_pos (int): Global number of true positives in target dataset

    Returns:
        pl.DataFrame: Dataset containing (`score_min`,`score_max`,`n_pos`,`n_neg`,`n`)
    """

    # derive score bin to observation bin mapping
    df_scorebins = (
        df_roc.sort("fpr")
        .with_columns(
            score_max=pl.col("fpr").rank("ordinal", descending=True)
            / (pl.col("fpr").len() - 1),
            cum_n_neg=pl.col("fpr").mul(n_neg).round(0).cast(pl.Int64),
            cum_n_pos=pl.col("tpr").mul(n_pos).round(0).cast(pl.Int64),
        )
        .with_columns(
            score_min=pl.col("score_max").shift(-1).fill_null(0),
            n_neg=pl.col("cum_n_neg") - pl.col("cum_n_neg").shift(1).fill_null(0),
            n_pos=pl.col("cum_n_pos") - pl.col("cum_n_pos").shift(1).fill_null(0),
        )
        .filter(pl.col("score_max").is_between(0, 1))
        .with_columns(n=pl.col("n_neg") + pl.col("n_pos"))
        .select("score_min", "score_max", "n_pos", "n_neg", "n")
    )

    return df_scorebins


def _gen_scorebins_to_scores(
    df_scorebins: pl.DataFrame
) -> pl.DataFrame:
    """Generate score-level dataset from scorebins

    Args:
        df_scorebins (pl.DataFrame): Scorebin data as returned by `_gen_roc_to_scorebins()`

    Returns:
        pl.DataFrame: Dataset containing `score` and `target` for each sample observation
    """

    def _score_fn(z):
        vals = np.linspace(z["score_min"],z["score_max"],z["n"])
        return np.sort(vals).tolist()

    def _target_fn(z):
        vals = np.append( 
                    np.repeat(0, z["n_neg"]), 
                    np.repeat(1, z["n_pos"])
                        )
        return np.sort(vals).tolist()

    df_scores = (
        df_scorebins
        .filter( pl.col('n') > 0)
        .with_columns(
            score=pl.struct("score_min", "score_max", "n").map_elements(
                function=_score_fn,
                return_dtype=pl.List(pl.Float64),
            ),
            target=pl.struct("n_pos", "n_neg").map_elements(
                function=_target_fn,
                return_dtype=pl.List(pl.Int64),
            ),
        )
        .explode("score", "target")
        .select("score", "target")
    )
    return df_scores
