import pandas as pd

def unify_author_fullname_and_orcid_core(
    df,
    orcid_col='Orcid',
    fullname_col='AuthorFullName',
    no_orcid_value='NO ORCID'
):
    """
    The original core function that does two steps:
      1) unify name by ORCID (pick the longest name),
      2) unify ORCID by the newly unified name (pick the most frequent ORCID, ignoring 'NO ORCID').

    Returns a DataFrame with 'UnifiedName' and 'UnifiedOrcid'.
    """
    df['Orcid'] = df['Orcid'].fillna('NO ORCID')# Remove '.' at the end of strings in a specific column (e.g., 'ColumnName')
    df['AuthorFullName'] = df['AuthorFullName'].str.rstrip('.')
    
    def pick_longest_name(names_series):
        freq = names_series.value_counts(dropna=False)
        freq_df = freq.reset_index()
        freq_df.columns = ['name', 'count']
        freq_df['length'] = freq_df['name'].str.len()

        # Sort by 'length' descending, then 'count' descending
        freq_df = freq_df.sort_values(by=['length', 'count'], ascending=[False, False])
        return freq_df['name'].iloc[0]

    # 1) Group by ORCID → unify name
    best_name_per_orcid = (
        df.groupby(orcid_col)[fullname_col].agg(pick_longest_name)
    )
    df['UnifiedName'] = df[orcid_col].map(best_name_per_orcid)

    def pick_most_frequent_orcid(orcids_series):
        freq = orcids_series.value_counts(dropna=False)
        if no_orcid_value in freq.index:
            freq = freq.drop(no_orcid_value)

        if freq.empty:
            return no_orcid_value
        else:
            return freq.idxmax()

    # 2) Group by UnifiedName → unify ORCID
    best_orcid_per_name = (
        df.groupby('UnifiedName')[orcid_col].agg(pick_most_frequent_orcid)
    )
    df['UnifiedOrcid'] = df['UnifiedName'].map(best_orcid_per_name)

    return df


def unify_author_fullname_and_orcid(
    wos_authors_enriched: pd.DataFrame,
    authorname_col='AuthorName',
    fullname_col='AuthorFullName',
    orcid_col='Orcid',
    max_unique_fullnames=5
):
    """
    Wrapper function that:
      - Finds each distinct AuthorName (except 'ANONYMOUS').
      - For each AuthorName, checks how many distinct AuthorFullName.
      - If <= max_unique_fullnames, applies the normal unification logic.
      - If > max_unique_fullnames, applies a fallback approach:
          * e.g. 'UnifiedName' = the original AuthorFullName
          * e.g. 'UnifiedOrcid' = the original ORCID

    Finally, concatenates all results together and returns a single DataFrame
    with new columns 'UnifiedName' and 'UnifiedOrcid'.
    """

    # We will accumulate results in a list, then concat them at the end
    all_results = []

    # 1) Identify all unique authorName values, sorted by frequency (descending) if you want
    # Exclude 'ANONYMOUS' from the loop
    authorname_counts = wos_authors_enriched[authorname_col].value_counts()
    unique_authornames = authorname_counts.index.tolist()

    for short_name in unique_authornames:
        # If it's ANONYMOUS, skip normal logic but keep them as-is
        if short_name.upper() == 'ANONYMOUS':
            # Just copy them out
            df_tmp = wos_authors_enriched[
                wos_authors_enriched[authorname_col] == short_name
            ].copy()

            # For ANONYMOUS, let's just keep them as-is.
            # or set the UnifiedName = AuthorFullName
            df_tmp['UnifiedName'] = df_tmp[fullname_col]
            df_tmp['UnifiedOrcid'] = df_tmp[orcid_col]
            all_results.append(df_tmp)
            continue

        # Filter the subset
        subset = wos_authors_enriched[
            wos_authors_enriched[authorname_col] == short_name
        ].copy()

        # 2) Count how many distinct AuthorFullName
        num_distinct = subset[fullname_col].nunique()

        # 3) Decide on the approach
        if num_distinct <= max_unique_fullnames:
            # Use the normal unify logic
            subset_result = unify_author_fullname_and_orcid_core(
                df=subset,
                orcid_col=orcid_col,
                fullname_col=fullname_col
            )
        else:
            # fallback approach
            # e.g. no merging: for each row, 'UnifiedName' = 'AuthorFullName'
            # and 'UnifiedOrcid' = 'Orcid'
            subset_result = subset.copy()
            subset_result['UnifiedName'] = subset_result[fullname_col]
            subset_result['UnifiedOrcid'] = subset_result[orcid_col]

        all_results.append(subset_result)

    # 4) Combine all subsets
    final_df = pd.concat(all_results, ignore_index=True)
    final_df['UnifiedOrcid'] = final_df['UnifiedOrcid'].fillna('NO ORCID')
    final_df['UnifiedName'] = final_df['UnifiedName'].str.rstrip('.')

    # Build a stable AuthorID preferring ORCID, then OpenAlexAuthorID, else name-based
    def _norm_orcid(val: str) -> str:
        if not isinstance(val, str):
            return ''
        s = val.strip()
        if not s or s.upper() == 'NO ORCID':
            return ''
        # Take last path segment if a URL, otherwise the raw string
        return s.split('/')[-1]

    def _short_openalex_id(val: str) -> str:
        if not isinstance(val, str):
            return ''
        s = val.strip()
        return s.split('/')[-1] if s else ''

    final_df['__OrcidNorm'] = final_df['UnifiedOrcid'].apply(_norm_orcid)
    if 'OpenAlexAuthorID' in final_df.columns:
        final_df['__OAshort'] = final_df['OpenAlexAuthorID'].apply(_short_openalex_id)
    else:
        final_df['__OAshort'] = ''

    # If an ORCID maps to multiple distinct OpenAlex IDs, prefer OA for disambiguation
    if 'OpenAlexAuthorID' in final_df.columns:
        # Build mapping: orcid_norm -> number of distinct OA ids (excluding empty)
        counts = (
            final_df.assign(__oa_nonempty=final_df['__OAshort'].astype(str).str.len() > 0)
                    .groupby('__OrcidNorm')['__OAshort']
                    .nunique(dropna=True)
        )
        # Convert to dict for fast lookup
        orcid_to_oa_unique = counts.to_dict()
    else:
        orcid_to_oa_unique = {}

    def _make_author_id(row):
        orcid_norm = row['__OrcidNorm']
        oa_short = row['__OAshort']
        # If this ORCID appears with multiple OpenAlex IDs, prefer OA to avoid collapsing distinct authors
        if orcid_norm and orcid_to_oa_unique.get(orcid_norm, 0) > 1 and oa_short:
            return f"OA:{oa_short}"
        if orcid_norm:
            return f"ORCID:{orcid_norm}"
        if oa_short:
            return f"OA:{oa_short}"
        # Fallback to name-based key (stable but less precise)
        base = f"{row.get('UnifiedName','') or ''}".strip()
        base = base.replace(',', '').replace('.', '')
        base = '_'.join(base.split())
        return f"NAME:{base.upper()}"

    final_df['AuthorID'] = final_df.apply(_make_author_id, axis=1)
    # Cleanup temp cols
    final_df = final_df.drop(columns=['__OrcidNorm', '__OAshort'], errors='ignore')

    final_df = final_df.rename(columns={
    'AuthorFullName': 'AuthorFullName_old',
    'UnifiedName': 'AuthorFullName',
    'Orcid': 'Orcid_old',
    'UnifiedOrcid': 'Orcid'
    })

    # Create the three DataFrames (carry OpenAlex IDs when available)
    wos_author_cols = [
        'AuthorID', 'AuthorName', 'AuthorFullName', 'Orcid', 'ResearcherID', 'Email',
        'OpenAlexAuthorID'
    ]
    wos_author = final_df[[c for c in wos_author_cols if c in final_df.columns]]
    wos_author = wos_author.drop_duplicates(subset=['AuthorID'])

    articleauthor_cols = [
        'SR', 'openalex_work_id', 'AuthorID', 'AuthorOrder', 'CorrespondingAuthor', 'OpenAlexAuthorID'
    ]
    articleauthor = final_df[[c for c in articleauthor_cols if c in final_df.columns]]
    articleauthor = articleauthor.drop_duplicates()

    wos_author_affiliation_cols = ['SR', 'AuthorID', 'Affiliation']
    wos_author_affiliation = final_df[[c for c in wos_author_affiliation_cols if c in final_df.columns]]
    wos_author_affiliation = wos_author_affiliation.drop_duplicates()

    # Return the three DataFrames as a list
    return wos_author, articleauthor, wos_author_affiliation
