import logging
from contextlib import contextmanager, nullcontext
from pathlib import Path

import bibtexparser
import pandas as pd


@contextmanager
def all_logging_disabled(highest_level=logging.CRITICAL):
    previous_level = logging.root.manager.disable
    logging.disable(highest_level)
    try:
        yield
    finally:
        logging.disable(previous_level)


def bib_library_from_dir(
    dir: Path, disable_warnings: bool = True
) -> bibtexparser.Library:
    bib_string = ""
    for partial_bib in dir.glob("**/*.bib"):
        with open(partial_bib) as f:
            bib_string += "\n".join(f.readlines())
    if disable_warnings:
        cm = all_logging_disabled
    else:
        cm = nullcontext
    with cm():
        return bibtexparser.parse_string(bib_string)


def bib_metadata_df(sample: bibtexparser.Library) -> pd.DataFrame:
    """Returns dataframe with relevant metadata extracted from a bibtex library.

    DataFrame contains: doi, the number times it has been cited in the
    literature (WOS statistic), its overall usage, and all keywords it
    has been assigned (in the reference management software).
    """
    return (
        pd.DataFrame(
            [
                [
                    entry["doi"] if "doi" in entry.fields_dict else None,
                    entry["times-cited"]
                    if "times-cited" in entry.fields_dict
                    else None,
                    entry["usage"] if "usage" in entry.fields_dict else None,
                    entry["keywords"] if "keywords" in entry.fields_dict else None,
                ]
                for entry in sample.entries
            ],
            columns=["doi", "cited", "usage", "keywords"],
        )
        .drop_duplicates("doi")
        .set_index("doi")
    )


def country_groups_df(wb_xlsx_file: Path) -> pd.DataFrame:
    """Returns a dataframe of WB country group definitions.

    DataFrame contains the region of a country, the country code,
    its income group, economy group and lending category.

    Contains an additional entry for 'global' with the country
    code 'WLD' which can be used for non-country specific
    observations (e.g. wide comparative studies using 30+
    developing/developed country datasets).
    """
    return pd.concat(
        [
            pd.read_excel(wb_xlsx_file),
            pd.DataFrame(
                data={
                    "Economy": ["global"],
                    "Code": ["WLD"],
                    "Region": [
                        "Europe & Central Asia;Sub-Saharan Africa;Latin America & Caribbean;East Asia & Pacific;Middle East & North Africa;South Asia;North America"
                    ],
                    "Income group": [""],
                    "Lending category": [""],
                }
            ),
        ]
    ).set_index("Economy")


def observations_with_metadata_df(
    raw_observations: pd.DataFrame,
    study_metadata: pd.DataFrame | None = None,
    country_groups: pd.DataFrame | None = None,
) -> pd.DataFrame:
    """Returns observations with metadata and country information attached.

    Returned DataFrame contains the raw observations data with one
    observation per row, with correct study and WB country grouping metadata
    attached for each row as well.
    """
    df = raw_observations.assign(
        doi=lambda _df: _df["uri"].str.extract(
            r"https?://(?:dx\.)?doi\.org/(.*)", expand=False
        ),
        date=lambda _df: pd.to_datetime(_df["year"], format="%Y"),
        year=lambda _df: _df["date"].dt.year,
    ).query("year >= 2000")
    if study_metadata is not None:
        df = df.assign(
            zot_cited=lambda _df: _df["doi"].map(study_metadata["cited"]),
            zot_usage=lambda _df: _df["doi"].map(study_metadata["usage"]),
            zot_keywords=lambda _df: _df["doi"].map(study_metadata["keywords"]),
        )
    if country_groups is not None:

        def c_to_region_closure(countries):
            return countries_to_regions(country_groups, countries)

        def c_to_income_closure(countries):
            return countries_to_income_groups(country_groups, countries)

        df = df.assign(
            region=lambda _df: _df["country"].map(c_to_region_closure),
            income_group=lambda _df: _df["country"].map(c_to_income_closure),
        )
    return df


def countries_to_regions(country_groups: pd.DataFrame, countries: str) -> str:
    """Returns the regions for countries passed in.

    Return is in semicolon-concatenated string-list ready to be passed into
    pandas for an apply or mapping method.
    """
    res = set()
    for c in countries.replace(" ;", ";").replace("; ", ";").split(";"):
        if c in country_groups.index:
            region = country_groups.at[c, "Region"]
            res.add(region)
    return ";".join(res)


def countries_to_income_groups(country_groups: pd.DataFrame, countries: str):
    """Returns the income groups for countries passed in.

    Return is in semicolon-concatenated string-list ready to be passed into
    pandas for an apply or mapping method.
    """
    res = set()
    for c in countries.replace(" ;", ";").replace("; ", ";").split(";"):
        if c in country_groups.index:
            region = country_groups.at[c, "Income group"]
            res.add(region)
    return ";".join(res)
