import math
from typing import cast

from pandas import DataFrame

METHOD_RANKINGS = {
    2.0: ["OLS", "ordinary.least.square", "logistic.regression", "fixed.effect"],
    3.0: [
        "DM",
        "discontinuity.matching",
        "DID",
        "difference.in.diff",
        "diff.in.diff",
        "triple.diff",
    ],
    3.5: ["PSM", "propensity.score.matching", "score.matching"],
    4.0: [
        "IV",
        "instrumental.variable",
        "method.of.moment",
        "GMM",
        "GEE",
        "generali(?:s|z)ed.estimating",
    ],
    4.5: ["RD", "regression.discontinuity"],
    5.0: ["RCT", "randomi(?:s|z)ed.control.trial"],
}


def calculate(
    df: DataFrame,
    repr_col: str = "representativeness",
    method_col: str = "method",
    valid_design_col: str | None = None,
) -> DataFrame:
    """Add internal and external validities to a dataframe.

    Requires a dataframe containing a study or observation per row, with a
    single column describing the study design, method and representativeness
    each respectively.

    Takes a combination of study design (simulation/observational/
    quasi-experimental/experimental/..) and its method (OLS/DID/RD/...) to
    calculate an internal validity. The column to use can be passed in as
    method_col.

    Takes a study representativeness (local/subnational/national/regional/
    census) to calculate the external validity. The column to use can be
    optionally passed in as repr_col.

    Takes an optional valid_design_col to decide if it should only
    include quasi-experimental/experimental designs in the calculation.
    Pass in the column name of the design column to only include those
    in calculations. Leave as None to calculate for all studies passed in.
    """
    EXT_COL_NAME: str = "external_validity"
    INT_COL_NAME: str = "internal_validity"
    cols = {EXT_COL_NAME: 0.0, INT_COL_NAME: 0.0}

    if valid_design_col:
        vd = df[
            (df[valid_design_col] == "quasi-experimental")
            | (df[valid_design_col] == "experimental")
        ].copy()
    vd = df.assign(**cols)
    vd = cast(DataFrame, vd)

    vd[repr_col] = vd[repr_col].fillna("")
    vd[method_col] = vd[method_col].fillna("")
    # needs to check national before subnational and census, subnational before local
    vd.loc[
        vd[repr_col].str.contains("|".join(["national", "regional"])), EXT_COL_NAME
    ] = 4.0
    vd.loc[vd[repr_col].str.contains("census"), EXT_COL_NAME] = 5.0
    vd.loc[vd[repr_col].str.contains("subnational"), EXT_COL_NAME] = 3.0
    vd.loc[vd[repr_col].str.contains("local"), EXT_COL_NAME] = 2.0

    # needs to go lowest to highest in case of multiple mentioned approaches
    for rank, methods in METHOD_RANKINGS.items():
        vd.loc[
            vd[method_col].str.contains("|".join(methods)),
            INT_COL_NAME,
        ] = rank

    return vd


def add_to_findings(
    findings_df: DataFrame,
    studies_by_intervention: DataFrame,
    strength_bins: dict[float, str] | None = None,
) -> DataFrame:
    """Returns summary of findings with validities added.

    Requires a 'findings' dataframe with at least a 'citation' column,
    containing the bibtex keys of studies containing the findings
    (semicolon-separated without spaces).
    Then, another dataframe with a row per study is required containing
    the study's 'method' and 'representativeness', in columns
    named respectively.

    Returns the correct bin to put the validity in, using the
    optionally passed in bins dictionary or the default 0 (weak) ->
    5 (evidence) -> 10 (strong) bins.
    """
    valid_subset = (
        calculate(studies_by_intervention)[
            ["internal_validity", "external_validity", "citation"]
        ]
        .fillna(1.0)
        .drop_duplicates(subset=["citation"])  # type: ignore
        .sort_values("internal_validity")
    )

    def apply_internal(df):
        return _combined_validities(
            df, valid_subset, "internal_validity", strength_bins
        )

    def apply_external(df):
        return _combined_validities(
            df, valid_subset, "external_validity", strength_bins
        )

    findings_df["internal_validity"] = findings_df["studies"].apply(apply_internal)
    findings_df["external_validity"] = findings_df["studies"].apply(apply_external)
    return findings_df


DEFAULT_BINS = {
    0.0: r"\-",
    1.0: r"\+",
    1.5: r"\++",
}


def _combined_validities(
    apply_to: DataFrame,
    by_intervention: DataFrame,
    column: str = "internal_validity",
    strength_bins: dict[float, str] | None = None,
):
    if not isinstance(apply_to, str):
        return
    combined = 0.0
    for study in apply_to.split(";"):
        if study not in by_intervention["citation"].unique():
            print(
                f"WARNING: Findings table {study} study did not match any study in interventions dataframe!"
            )
        new = by_intervention.loc[by_intervention["citation"] == study, column]
        if len(new) == 0 or math.isnan(new.iat[0]):
            continue
        combined += new.iat[0]

    if combined:
        return _binned_strength(combined, bins=strength_bins or DEFAULT_BINS)
    return r"\-"


def _binned_strength(strength: float, bins: dict[float, str]) -> str:
    bin = ""
    for val, txt in sorted(bins.items()):
        if strength >= val:
            bin = txt
    return bin
