import itertools
import json
import hashlib
import os
from collections import Counter
from datetime import datetime
from pathlib import Path

import nomquamgender as nqg
import gradio as gr
import polars as pl
import Levenshtein
import plotly.express as px
import plotly.graph_objects as go
import pyalex
import requests
import spaces
from grobid_client.grobid_client import GrobidClient
from pyalex import Works, Authors


OPENALEX_API_KEY = os.getenv('OA_KEY')
pyalex.config.api_key = OPENALEX_API_KEY



@spaces.GPU
def _null():
    pass

def hash_file(filepath: str|Path) -> str:
    BUF_SIZE = 65536
    sha = hashlib.sha256()

    with open(filepath, 'rb') as f:
        while True:
            data = f.read(BUF_SIZE)
            if not data:
                break
            sha.update(data)
    return sha.hexdigest()


def get_all_fields() -> list[str]:
    return pl.read_parquet('topics.parquet')['subfield'].unique().sort().to_list()


def parse_and_get_references(file: Path) -> list[dict]:
    client = GrobidClient('https://grobidorg-grobid.hf.space')
    client.process(
        "processReferences",
        file.parent,
        json_output=True
    )
    with file.with_suffix('.json').open() as f:
        refs = json.load(f)['references']
    refs = {ref['title'] for ref in refs if 'title' in ref}
    refs = [
        {'title': ref, 'openalex': None}
        for ref in refs
    ]
    for ref in refs:
        ref_title = ''.join(c if c.isalnum() else ' ' for c in ref['title'])
        works = Works().search_filter(title=ref_title).get(per_page=1)
        if len(works) == 0:
            print(f'no works for {ref_title}')
            continue
        alex_title = works[0]['title']
        if Levenshtein.ratio(alex_title.lower(), ref_title.lower()) > 0.9:
            ref['openalex'] = works[0]
        elif len(works) > 0:
            score = Levenshtein.ratio(alex_title.lower(), ref_title.lower())
            print(f'no paper match ({score}) between "{alex_title}" and "{ref_title}"')
    return refs


def get_gender(model: nqg.NBGC, name: str) -> str:
    if ',' in name:
        lastpart, firstpart = name.split(',', 1)
        name = firstpart.strip() + ' ' + lastpart.strip()
    guess = model.classify(name)[0]
    if guess == 'gm':
        return 'men'
    elif guess == 'gf':
        return 'women'
    else:
        print('guesser -', name)
        return 'unknown'


def calculate_gender_statistics(df: pl.DataFrame) -> tuple[pl.DataFrame, float]:
    #uroman = ur.Uroman()
    model = nqg.NBGC()
    gender_df = (df
        .filter(pl.col('openalex').struct.field('authorships').list.len() > 0)
        .with_columns(pl.col('openalex')
            .struct.field('authorships')
            .list.get(0)
            .struct.field('author')
            .struct.field('display_name')
            .map_elements(lambda name: get_gender(model, name))
            .alias('gender')
        )
        .select('title', 'gender')
    )
    df = df.join(gender_df, on='title', how='left')

    count_men = (df['gender'] == 'men').sum()
    count_women = (df['gender'] == 'women').sum()

    return df, count_women / max(count_men + count_women, 1)


def get_gn_gs_for_country_code(cc: str) -> str:
    try:
        return (
            pl.read_csv('global_north_south.csv')
            .filter(pl.col('country_code') == cc)
            ['Region'][0]
        )
    except IndexError:
        print('Could not find country', cc)
        return 'unknown'


def calculate_north_south_statistics(df: pl.DataFrame) -> tuple[pl.DataFrame, float]:
    region_df = (df
        .with_columns(pl.col("openalex").struct.field("authorships"))
        .filter(pl.col("authorships").list.len() > 0)
        .filter(pl.col("authorships").list.get(0).struct.field("countries").list.len() > 0)
        .with_columns(pl.col('authorships')
            .list.get(0)
            .struct.field('countries')
            .list.get(0)
            .map_elements(get_gn_gs_for_country_code, return_dtype=pl.Utf8)
            .alias('region')
        )
        .drop('authorships')
        .select('title', 'region')
    )
    df = df.join(region_df, on='title', how='left')

    count_gs = (df['region'] == 'GS').sum()
    count_gn = (df['region'] == 'GN').sum()

    return df, count_gs / max(count_gs + count_gn, 1)


@gr.cache
def extract_file_header_info(file: Path) -> dict:
    client = GrobidClient('https://grobidorg-grobid.hf.space')
    client.process(
        "processHeaderDocument",
        file.parent,
        json_output=True
    )
    with file.with_suffix('.json').open() as f:
        header = json.load(f)['biblio']
    if len(header.get('title', '')) < 10:
        raise ValueError('Title not found or too short')
    print(header)
    return {
        'title': header.get('title', ''),
        'abstract': (header.get('abstract') or [{}])[0].get('text', ''),
        'authors': header.get('authors', [])
    }


def clean_string(text: str) -> str:
    return ''.join([
        c
        if c.isalnum() or c in [' ', '-', ':']
        else ' '
        for c in text
    ])


@gr.cache
def classify_work_subfield(file: Path) -> dict:
    fileinfo = extract_file_header_info(file)
    # title and abstract must not be longer than 2000 characters
    title = clean_string(fileinfo['title'])
    abstract = clean_string(fileinfo['abstract'])[:1999-len(title)]

    url = f"https://api.openalex.org/text/topics?api_key={OPENALEX_API_KEY}&title={title}&abstract={abstract}"
    response = requests.get(url)
    return response.json()['primary_topic']['subfield']


@gr.cache
def get_openalex_authors(file: Path) -> list[dict]:
    authors = extract_file_header_info(file)['authors']
    authors = {author: [] for author in authors}

    # get matches and collaborators from OpenAlex
    for author in authors.keys():
        matches = Authors().search(author).get()[:5]
#       matches.sort(key=lambda m: Levenshtein.distance(author, m['display_name']))
        for author_match in matches:
            query = Works().filter(author={'id': author_match['id']}).group_by('author.id')
            collaborators = list(itertools.chain(*query.paginate(per_page=200)))[1:]
            author_match['collaborators'] = {c['key'] for c in collaborators}
            authors[author].append(author_match)

    # use match with the highest collaborator overlap
    best_match_score = 0
    best_match_combo = []
    iter_list = [[None] if not matches else matches for matches in authors.values()]
    for author_combination in itertools.product(*iter_list):
        collab_multiple = Counter([
            collab
            for author in author_combination
            for collab in author['collaborators'] or []
        ])
        score = sum(c for c in collab_multiple.values() if c > 1)
        if score > best_match_score or len(best_match_combo) == 0:
            best_match_score = score
            best_match_combo = author_combination
    return [author for author in best_match_combo if author is not None]


@gr.cache
def calculate_selfcites_statistics(df: pl.DataFrame, file: Path) -> tuple[pl.DataFrame,float|None]:
    authors = get_openalex_authors(file)
    if len(authors) == 0:
        return df, None
    author_selfcites = {a['display_name'] for a in authors}
    count = (df
        .filter(pl.col('openalex').is_not_null())
        .with_columns(pl.col('openalex').struct.field('authorships'))
        .explode('authorships')
        .with_columns(pl.col('authorships')
            .struct.field('author')
            .struct.field('display_name')
            .is_in(author_selfcites)
            .alias('is_self_cite')
        )
        .group_by('title')
        .agg(pl.col('is_self_cite').any())
    )
    df = df.join(count, on='title', how='left')
    return df, df['is_self_cite'].sum() / (len(count) or 1)


def gender_stat_subfield(field: str) -> pl.DataFrame:
    return (
        pl.read_parquet('women_per_year_per_subfield.parquet')
        .filter(pl.col('subfield') == field)
        .rename({
            'publication_year': 'Publication Year',
            'women_ratio': 'Women Ratio'
        })
    )


def region_stat_subfield(field: str) -> pl.DataFrame:
    return (
        pl.read_parquet('region_per_year_per_subfield.parquet')
        .filter(pl.col('subfield') == field)
        .rename({
            'publication_year': 'Publication Year',
            'global_south_ratio': 'Global South Ratio',
            'Region': 'Citing Paper Region'
        })
    )


def selfcites_stat_subfield(field: str) -> pl.DataFrame:
    return (
        pl.read_parquet('selfcites_per_year_per_subfield.parquet')
        .filter(pl.col('subfield') == field)
        .rename({
            'publication_year': 'Publication Year',
            'selfcite_ratio': 'Self-Citations Ratio'
        })
    )


def update_dropdown(state: dict, selected_field: str):
    gender_ratio = gender_stat_subfield(selected_field)
    region_ratio = region_stat_subfield(selected_field)
    selfcites_ratio = selfcites_stat_subfield(selected_field)

    gender_chart = px.line(
        gender_ratio,
        x="Publication Year",
        y="Women Ratio",
        template="plotly_white",
        title='Share of Women as First Author in References'
    )
    gender_chart.update_layout(xaxis={'ticklabelstandoff': 10})

    if state is not None and state["gender_stats_paper"] is not None:
        gender_chart.update_yaxes(range=[
            0,
            max(gender_ratio["Women Ratio"].max() + 0.1, state['gender_stats_paper'] + 0.1)
        ])
        gender_chart.add_trace(
            go.Scatter(
                x=[datetime.now().year],
                y=[state["gender_stats_paper"]],
                mode="markers",
                marker=dict(
                    size=10,
                    color="red",
                    line=dict(width=2, color='DarkSlateGrey')
                ),
                name="Uploaded Paper"
            )
        )

    region_chart = px.line(
        region_ratio,
        x="Publication Year",
        y="Global South Ratio",
        color="Citing Paper Region",
        template="plotly_white",
        title='Share of Global South Affiliation for First Author in References'
    )
    region_chart.update_layout(xaxis={'ticklabelstandoff': 10})
    if state is not None and state["region_stats_paper"] is not None:
        region_chart.update_yaxes(range=[
            0,
            max(region_ratio["Global South Ratio"].max() + 0.1, state['region_stats_paper'] + 0.1)
        ])
        region_chart.add_trace(
            go.Scatter(
                x=[datetime.now().year],
                y=[state["region_stats_paper"]],
                mode="markers",
                marker=dict(
                    size=10,
                    color="red",
                    line=dict(width=2, color='DarkSlateGrey')
                ),
                name="Uploaded Paper"
            )
        )

    selfcites_chart = px.line(
        selfcites_ratio,
        x="Publication Year",
        y="Self-Citations Ratio",
        template="plotly_white",
        title='Share of Self-Citations'
    )
    selfcites_chart.update_layout(xaxis={'ticklabelstandoff': 10})
    if state is not None and state["selfcites_stats_paper"] is not None:
        selfcites_chart.update_yaxes(range=[
            0,
            max(selfcites_ratio["Self-Citations Ratio"].max() + 0.1, state['selfcites_stats_paper'] + 0.1)
        ])
        selfcites_chart.add_trace(
            go.Scatter(
                x=[datetime.now().year],
                y=[state["selfcites_stats_paper"]],
                mode="markers",
                marker=dict(
                    size=10,
                    color="red",
                    line=dict(width=2, color='DarkSlateGrey')
                ),
                name="Uploaded Paper"
            )
        )

    return gender_chart, region_chart, selfcites_chart


@gr.cache(key=lambda args: hash_file(args['file']))
async def upload_file(file: str):
    try:
        pfile = Path(file)

        subfield = classify_work_subfield(pfile)['display_name']

        refs = parse_and_get_references(pfile)
        refs_df = pl.DataFrame(refs)
        refs_df, gender_stats_paper = calculate_gender_statistics(refs_df)
        refs_df, region_stats_paper = calculate_north_south_statistics(refs_df)
        refs_df, selfcites_stats_paper = calculate_selfcites_statistics(refs_df, pfile)

        refs_df = (
            refs_df
            .with_columns(pl.col('openalex').struct.field('id').alias('OpenAlex ID'))
            .rename({'title': 'Title', 'gender': 'Inferred Gender', 'region': 'Region', 'is_self_cite': 'Self-citation'})
        )
        styled_df = refs_df.drop('openalex').to_pandas()

        authors_names = [f"- {a['id']} {a['display_name']}" for a in get_openalex_authors(pfile)]

        def percent_or_none(n: float|None) -> str:
            if n is None:
                return 'unable to calculate'
            return f'{round(n * 100, 2)}%'

        result_text = (f'Detected field: {subfield}.\n\n'
            f'Matched OpenAlex authors:\n{'\n'.join(authors_names)}\n\n'
            f'Global south ratio: {percent_or_none(region_stats_paper)}\n'
            f'Women ratio:  {percent_or_none(gender_stats_paper)}\n'
            f'Self-citations ratio: {percent_or_none(selfcites_stats_paper)}'
        )

        state = {
            "field": subfield,
            "gender_stats_paper": gender_stats_paper,
            "region_stats_paper": region_stats_paper,
            "selfcites_stats_paper": selfcites_stats_paper,
        }

        return state, file, result_text, styled_df, subfield
    except Exception as e:
        error_msg = f'Failed with error: {e}'
        return {}, file, error_msg, None, None


if __name__ == '__main__':
    css = """
    .gradio-container {
        max-width: 1000px !important;
        margin: auto !important;
    }
    """

    with gr.Blocks(css=css, title="Reference Analysis Prototype") as demo:
        state = gr.State()
        gr.Markdown("""
        # CiteCheck

        Upload a PDF to analyze its references. The app extracts references and authors,
        matches them to OpenAlex, and computes statistics in comparison the article's research
        field.

        The website computes:
        - the number of references to women as first authors (based on inferring gender from name)
        - the number of references to the Global South for first author (based on the affiliation); data shown separately for citing articles from Global North and Global South since local research is automatically referenced more
        - the number of self-citations in the references

        **These numbers can be inaccurate for individual papers, so use them with caution.**
        They are meant to raise awareness on citation behaviour and not as exact metrics.
        """)
        with gr.Row():
            with gr.Column():
                file_output = gr.File()
                upload_button = gr.UploadButton(
                    "Click to Upload a File",
                    file_types=[".pdf"]
                )
                field_drop = gr.Dropdown(choices=get_all_fields(), label='Field')
            with gr.Column():
                result_out = gr.Textbox(interactive=False)
        with gr.Row():
            gender_chart = gr.Plot()
        with gr.Row():
            region_chart = gr.Plot()
        with gr.Row():
            selfcites_chart = gr.Plot()
        with gr.Row():
            result_table_paper = gr.Dataframe()
        field_drop.change(
            update_dropdown,
            inputs=[state, field_drop],
            outputs=[gender_chart, region_chart, selfcites_chart]
        )
        upload_button.upload(
            upload_file,
            upload_button,
            [
                state,
                file_output,
                result_out,
                result_table_paper,
                field_drop
             ]
        )

    demo.launch(max_file_size="10mb")