"""
run_harmonization.py

Reads all standardized JSON files from fetch_prwp, aggregates raw dataset
mention frequencies, and runs the ai4data harmonization pipeline to produce
a canonical_map.json lookup: raw_variant_text -> formal canonical name.

Optimized to run clustering and country lookup in O(N) vectorized logic,
and configured to run on MPS (GPU) on macOS.
"""
import glob
import json
import os
import sys
import re
from collections import Counter
from pathlib import Path

import nltk
import numpy as np
import pandas as pd
from tqdm.auto import tqdm

# Ensure ai4data is importable
AI4DATA_SRC = "/Users/rafaelmacalaba/WBG/ai4data/src"
if AI4DATA_SRC not in sys.path:
    sys.path.insert(0, AI4DATA_SRC)

# ──────────────────────────────────────────────────────────────────────────────
# Monkey-patch fast country detection and fast clustering
# ──────────────────────────────────────────────────────────────────────────────
print("Initializing environment and pre-compiling country/city lookup regex...")
import ai4data.data_use.extractors.harmonization as harm

# Load custom country_map.json containing demonyms/adjectives
country_map_path = Path("/Users/rafaelmacalaba/WBG/ai4data/src/ai4data/data_use/assets/country_map.json")
with open(country_map_path, "r", encoding="utf-8") as f:
    user_country_map = json.load(f)

country_map = harm.build_country_map_with_cities_only(user_country_map)
form_to_country = {}
for country, forms in country_map.items():
    for f in forms:
        form_to_country[harm.normalize(f)] = country

# Sort longest first so multi-word forms match first
sorted_forms = sorted(form_to_country.keys(), key=len, reverse=True)
country_detection_regex = re.compile(
    r"\b(" + "|".join(map(re.escape, sorted_forms)) + r")\b"
)

def detect_country_fast(raw: str, country_map_ignored=None) -> str | None:
    if not isinstance(raw, str):
        return None
    clean = harm.normalize(raw)
    match = country_detection_regex.search(clean)
    if match:
        return form_to_country.get(match.group(1))
    return None

# Override the slow nested-loop implementation with the optimized regex
harm.detect_country = detect_country_fast
print("Monkey-patched detect_country successfully.")

def learn_family_keys_safe(families, sim_threshold=85, sem_threshold=0.8):
    """
    Safe version of learn_family_keys that guards acronym check to prevent
    AttributeError: 'float' object has no attribute 'lower' when acronym is np.nan.
    Resolves acronym conflicts by keeping the family with the highest mention count.
    """
    family_keys = {}
    acronym_best = {} # acr_lower -> (canonical_name, total_count)

    for fam in families:
        cname = fam["Canonical"]["raw_name"]
        base_norm = fam["Canonical"].get("base_name_norm", cname.lower())
        acr = fam["Canonical"].get("acronym")

        canonical_name = cname
        if acr and isinstance(acr, str) and acr.strip():
            canonical_name = f"{cname} ({acr})"

        variants = []
        counts = Counter()

        # Calculate total family count to resolve acronym conflicts
        fam_count = fam["Canonical"].get("count", 1)

        # Add aliases
        for alias in fam.get("Aliases", []):
            norm = alias.get("base_name_norm", alias["raw_name"].lower())
            variants.append(norm)
            c = alias.get("count", 1)
            counts[norm] += c
            fam_count += c

        # Add prototypes and their aliases
        for proto in fam.get("Prototypes", []):
            pnorm = proto["Prototype"].get("base_name_norm", proto["Prototype"]["raw_name"].lower())
            variants.append(pnorm)
            c = proto["Prototype"].get("count", 1)
            counts[pnorm] += c
            fam_count += c

            for a in proto.get("Aliases", []):
                anorm = a.get("base_name_norm", a["raw_name"].lower())
                variants.append(anorm)
                c_a = a.get("count", 1)
                counts[anorm] += c_a
                fam_count += c_a

        family_keys[base_norm] = canonical_name

        for v in set(variants):
            if acr and isinstance(acr, str) and harm.is_acronym_variant(v, acr):
                family_keys[v] = canonical_name
            else:
                match = harm.process.extractOne(v, [base_norm], scorer=harm.fuzz.ratio)
                if match and match[1] >= sim_threshold:
                    family_keys[v] = canonical_name

        if acr and isinstance(acr, str) and acr.strip():
            acr_key = acr.lower()
            if acr_key in acronym_best:
                prev_name, prev_count = acronym_best[acr_key]
                if fam_count > prev_count:
                    acronym_best[acr_key] = (canonical_name, fam_count)
            else:
                acronym_best[acr_key] = (canonical_name, fam_count)

    # Apply the best (highest frequency) acronym mappings
    for acr_key, (canonical_name, _) in acronym_best.items():
        family_keys[acr_key] = canonical_name

    return family_keys

harm.learn_family_keys = learn_family_keys_safe
print("Monkey-patched learn_family_keys successfully.")

def merge_acronyms_safe(families, sim_threshold=0.8):
    """
    Safe version of merge_acronyms that checks isinstance(acr, str)
    to prevent AttributeError: 'float' object has no attribute 'lower'.
    """
    merged = []
    used = set()

    for i, fam in enumerate(families):
        if i in used:
            continue

        canonical = fam["Canonical"]
        acr = canonical.get("acronym")

        if acr and isinstance(acr, str) and acr.strip():
            longform_family = fam

            for j, other in enumerate(families):
                if j == i or j in used:
                    continue

                other_name = other["Canonical"]["raw_name"]
                other_base = other["Canonical"].get("base_name_norm", other_name.lower())

                # Check if acronym is in the other canonical name
                if acr.lower() in other_name.lower() or acr.lower() in other_base:
                    longform_family["Aliases"].append(other["Canonical"])
                    longform_family["Aliases"].extend(other.get("Aliases", []))
                    longform_family["Prototypes"].extend(other.get("Prototypes", []))
                    used.add(j)

            merged.append(longform_family)
            used.add(i)
        else:
            merged.append(fam)
            used.add(i)

    return merged

harm.merge_acronyms = merge_acronyms_safe
print("Monkey-patched merge_acronyms successfully.")

def cluster_names_fast(df, embedder, sim_threshold=0.85):
    """
    Vectorized version of cluster_names that performs row-wise thresholding in numpy
    instead of nested loops in Python.
    """
    # Step 1: Pre-filter
    df_filtered = harm.prefilter(df).reset_index(drop=True)
    if df_filtered.empty:
        df_filtered["cluster"] = []
        return df_filtered

    # Step 2: Compute similarity matrix
    sim = harm.compute_hybrid_similarity(df_filtered, embedder)

    # Step 3: Fast vector clustering
    visited = set()
    cluster_labels = np.full(len(df_filtered), -1)
    cluster_id = 0

    for i in range(len(df_filtered)):
        if i in visited:
            continue

        # Vectorized check for similarity >= threshold in row i
        matching_indices = np.where(sim[i] >= sim_threshold)[0]

        cluster_idx = [i]
        visited.add(i)

        for j in matching_indices:
            if j > i and j not in visited:
                cluster_idx.append(j)
                visited.add(j)

        cluster_labels[cluster_idx] = cluster_id
        cluster_id += 1

    df_filtered["cluster"] = cluster_labels
    return df_filtered

harm.cluster_names = cluster_names_fast
print("Monkey-patched cluster_names successfully.")

# Import remaining harmonization functions
from ai4data.data_use.extractors.harmonization import (
    build_country_regex,
    build_families,
    learn_family_keys,
    normalize,
    preprocess_cluster,
    merge_acronyms,
    consolidate_families,
)

# ──────────────────────────────────────────────────────────────────────────────
# Configuration
# ──────────────────────────────────────────────────────────────────────────────
BASE_DIR = Path(__file__).parent
STANDARDIZED_BASE = Path("/Users/rafaelmacalaba/WBG/fetch_prwp/data/standardized_outputs")
OUTPUT_PATH = BASE_DIR / "canonical_map.json"
FAMILIES_OUTPUT_PATH = BASE_DIR / "dataset_families.json"

# Only process mentions where specificity_tag == "named" (formal, named datasets)
NAMED_ONLY = True

# Harmonization similarity threshold (from the original code's defaults)
SIM_THRESHOLD = 0.82


# ──────────────────────────────────────────────────────────────────────────────
# Step 1: Read all standardized JSONs and collect raw mentions
# ──────────────────────────────────────────────────────────────────────────────
def collect_raw_mentions(standardized_base: Path, named_only: bool = True) -> pd.DataFrame:
    all_json_files = glob.glob(str(standardized_base / "batch_*" / "*.json"))
    print(f"  Found {len(all_json_files)} standardized JSON files.")

    # Aggregate: raw_name -> { count, acronyms[] }
    name_counts: Counter = Counter()
    name_acronyms: dict[str, Counter] = {}

    for filepath in tqdm(all_json_files, desc="  Reading files", unit="file"):
        try:
            with open(filepath, "r", encoding="utf-8") as f:
                doc = json.load(f)
        except Exception:
            continue

        for extraction in doc.get("model_extractions") or []:
            if extraction.get("classifier_skipped", False):
                continue

            for ds in extraction.get("datasets") or []:
                specificity = (ds.get("specificity_tag") or {}).get("text", "").strip().lower()
                if named_only and specificity != "named":
                    continue

                mention = (ds.get("mention_name") or {}).get("text", "").strip()
                if not mention or len(mention) < 4:
                    continue

                acronym = (ds.get("acronym") or {}).get("text", "").strip()

                name_counts[mention] += 1
                if mention not in name_acronyms:
                    name_acronyms[mention] = Counter()
                if acronym:
                    name_acronyms[mention][acronym] += 1

    print(f"  Collected {len(name_counts)} unique raw mention strings.")

    # Build DataFrame
    rows = []
    for raw_name, count in name_counts.items():
        best_acronym = None
        if name_acronyms.get(raw_name):
            best_acronym = name_acronyms[raw_name].most_common(1)[0][0]
        rows.append({
            "raw_name": raw_name,
            "count": count,
            "acronym": best_acronym,
        })

    return pd.DataFrame(rows)


# ──────────────────────────────────────────────────────────────────────────────
# Step 2: Preprocess into base_name_norm using the harmonization utilities
# ──────────────────────────────────────────────────────────────────────────────
def preprocess_df(df: pd.DataFrame):
    country_pattern = build_country_regex(country_map)

    # Download required NLTK data silently
    nltk.download("stopwords", quiet=True)
    nltk.download("wordnet", quiet=True)
    from nltk.corpus import stopwords as nltk_stopwords
    from nltk.stem import WordNetLemmatizer

    lemmatizer = WordNetLemmatizer()
    stop_words = set(nltk_stopwords.words("english"))

    print("  Preprocessing raw names (stripping countries, years, normalizing)...")
    preprocessed = preprocess_cluster(
        df,
        country_map=country_map,
        country_pattern=country_pattern,
        lemmatizer=lemmatizer,
        stopwords=stop_words,
    )
    return preprocessed


# ──────────────────────────────────────────────────────────────────────────────
# Step 3: Run clustering and hierarchical family building
# ──────────────────────────────────────────────────────────────────────────────
def run_harmonization(preprocessed_df: pd.DataFrame):
    import torch
    device = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu")
    print(f"  Using device for sentence-transformer: {device}")

    print("  Loading sentence-transformer embedder...")
    from sentence_transformers import SentenceTransformer
    embedder = SentenceTransformer("all-MiniLM-L6-v2", device=device)

    # Filter out single-mention items to speed up execution
    # (keeps items with count >= 2, reducing unique names from 38k to 9k)
    print("  Filtering raw dataset mentions (keeping count >= 2 for clustering)...")
    df_frequent = preprocessed_df[preprocessed_df["count"] >= 2].reset_index(drop=True)
    print(f"  Reduced unique name pool to {len(df_frequent)} entries.")

    # Prepare DataFrame columns for cluster_names helper
    df_for_clustering = df_frequent.rename(columns={"raw_name": "datasets"})

    print(f"  Clustering names with hybrid similarity (sim_threshold={SIM_THRESHOLD})...")
    df_clustered = harm.cluster_names(
        df_for_clustering[["datasets", "count", "acronym", "base_name_norm", "country", "base_name"]],
        embedder,
        SIM_THRESHOLD
    )

    cluster_labels = df_clustered["cluster"].unique()
    print(f"  Found {len(cluster_labels)} similarity clusters.")

    # Replace any NaN/float-nulls with None to prevent downstream AttributeError: 'float' object has no attribute 'lower' in learn_family_keys
    df_clustered = df_clustered.where(pd.notna(df_clustered), None)

    import copy

    all_families = []
    all_unconsolidated_families = []

    for cluster_id in tqdm(cluster_labels, desc="  Building hierarchies per cluster"):
        df_batch = df_clustered[df_clustered["cluster"] == cluster_id].rename(
            columns={"datasets": "raw_name"}
        )
        
        # Build hierarchy for this cluster (preprocessed, no need to run preprocess_cluster again!)
        families = build_families(df_batch, sim_threshold=0.85)
        families = merge_acronyms(families)
        
        # Accumulate the unconsolidated families for global acronym and variant learning
        all_unconsolidated_families.extend(copy.deepcopy(families))
        
        family_keys = learn_family_keys(families, sim_threshold=85)
        families = consolidate_families(families, family_keys, sim_threshold=85)
        
        all_families.extend(families)

    print("  Learning global family keys on all unconsolidated families...")
    all_family_keys = learn_family_keys(all_unconsolidated_families, sim_threshold=85)

    # Map remaining single-mention names to the learned family keys where possible
    print("  Mapping single-mention names to learned canonical keys...")
    df_singles = preprocessed_df[preprocessed_df["count"] < 2].reset_index(drop=True)
    mapping_hits = 0
    for _, row in df_singles.iterrows():
        raw = row["raw_name"]
        norm = row["base_name_norm"]
        acronym = row["acronym"]
        
        # Check if normalized base name or acronym matches a canonical key
        matched_canonical = None
        acronym_str = acronym.lower() if isinstance(acronym, str) else ""
        for key in [norm, raw.lower(), acronym_str]:
            if key and key in all_family_keys:
                matched_canonical = all_family_keys[key]
                break
        
        if matched_canonical:
            all_family_keys[raw] = matched_canonical
            mapping_hits += 1

    print(f"  Mapped {mapping_hits} single-mention names to canonical families.")
    print(f"  Resolved {len(all_family_keys)} variant -> canonical mappings in total.")
    return all_families, all_family_keys


# ──────────────────────────────────────────────────────────────────────────────
# Step 4: Save outputs
# ──────────────────────────────────────────────────────────────────────────────
def save_outputs(families: list, family_keys: dict):
    os.makedirs(OUTPUT_PATH.parent, exist_ok=True)

    # Save the canonical_map: variant -> canonical_name
    with open(OUTPUT_PATH, "w", encoding="utf-8") as f:
        json.dump(family_keys, f, indent=2, ensure_ascii=False)
    print(f"  Saved canonical map to: {OUTPUT_PATH}")

    # Save the full families structure (for inspection/debugging)
    def make_serializable(obj):
        if isinstance(obj, dict):
            return {k: make_serializable(v) for k, v in obj.items()}
        elif isinstance(obj, list):
            return [make_serializable(i) for i in obj]
        elif isinstance(obj, (np.integer,)):
            return int(obj)
        elif isinstance(obj, (np.floating,)):
            return float(obj)
        elif isinstance(obj, float) and (obj != obj):  # NaN
            return None
        return obj

    with open(FAMILIES_OUTPUT_PATH, "w", encoding="utf-8") as f:
        json.dump(make_serializable(families), f, indent=2, ensure_ascii=False)
    print(f"  Saved full families to: {FAMILIES_OUTPUT_PATH}")


def consolidate_acronym_families(family_keys: dict, preprocessed_df) -> dict:
    """
    Consolidates variant/acronym mappings of major datasets into their primary canonical parents.
    Iterates over all known raw names and acronyms to ensure complete coverage.
    """
    import re
    
    # Define primary canonical targets
    dhs_target = "Demographic and Health Surveys (DHS)"
    lsms_target = "Living Standards Measurement Study (LSMS)"
    wdi_target = "World Development Indicators (WDI)"
    
    # Pre-populate with existing mappings
    consolidated = {}
    for variant, canonical in family_keys.items():
        consolidated[variant] = canonical
        
    # We will check all raw names in the preprocessed pool to be 100% comprehensive
    unique_raw_names = preprocessed_df["raw_name"].unique()
    
    for raw in unique_raw_names:
        r_lower = raw.lower().strip()
        
        # Regex patterns for matching
        # DHS patterns:
        is_dhs = (
            "demographic and health" in r_lower
            or "demographic and heath" in r_lower
            or "demographic & health" in r_lower
            or "demographic & heath" in r_lower
            or "demographic and household" in r_lower
            or "demographic & household" in r_lower
            or "demographic health survey" in r_lower
            or "demographic heath survey" in r_lower
            or re.search(r"\b[a-z]?dhs\b", r_lower) is not None
        ) and not any(x in r_lower for x in ["cdhs", "ais", "asset index", "dhs/ais", "dhs/cov"])
        
        # LSMS patterns:
        is_lsms = (
            "living standards" in r_lower
            or "living standard" in r_lower
            or re.search(r"\blsms\b", r_lower) is not None
        )
        
        # WDI patterns:
        is_wdi = (
            "world development" in r_lower
            or re.search(r"\bwdi\b", r_lower) is not None
        ) and not any(x in r_lower for x in ["wvs", "sarmd", "economic freedom", "world development report", "wdr"])
        
        if is_dhs:
            consolidated[raw] = dhs_target
            consolidated[r_lower] = dhs_target
        elif is_lsms:
            consolidated[raw] = lsms_target
            consolidated[r_lower] = lsms_target
        elif is_wdi:
            consolidated[raw] = wdi_target
            consolidated[r_lower] = wdi_target
            
    # Also apply the same rules to update existing keys in consolidated
    for variant, canonical in list(consolidated.items()):
        v_lower = variant.lower()
        c_lower = canonical.lower()
        
        is_dhs = (
            "demographic and health" in v_lower
            or "demographic and heath" in v_lower
            or "demographic & health" in v_lower
            or "demographic & heath" in v_lower
            or "demographic and household" in v_lower
            or "demographic & household" in v_lower
            or "demographic health survey" in v_lower
            or "demographic heath survey" in v_lower
            or re.search(r"\b[a-z]?dhs\b", v_lower) is not None
            or "demographic and health" in c_lower
            or "demographic and heath" in c_lower
            or "demographic & health" in c_lower
            or "demographic & heath" in c_lower
            or "demographic and household" in c_lower
            or "demographic & household" in c_lower
            or "demographic health survey" in c_lower
            or "demographic heath survey" in c_lower
            or re.search(r"\b[a-z]?dhs\b", c_lower) is not None
        ) and not any(x in v_lower for x in ["cdhs", "ais", "asset index", "dhs/ais", "dhs/cov"]) \
          and not any(x in c_lower for x in ["cdhs", "ais", "asset index", "dhs/ais", "dhs/cov"])
          
        is_lsms = (
            "living standards" in v_lower
            or "living standard" in v_lower
            or re.search(r"\blsms\b", v_lower) is not None
            or "living standards" in c_lower
            or "living standard" in c_lower
            or re.search(r"\blsms\b", c_lower) is not None
        )
        
        is_wdi = (
            ("world development" in v_lower or re.search(r"\bwdi\b", v_lower) is not None)
            and not any(x in v_lower for x in ["wvs", "sarmd", "economic freedom", "world development report", "wdr"])
        ) or (
            ("world development" in c_lower or re.search(r"\bwdi\b", c_lower) is not None)
            and not any(x in c_lower for x in ["wvs", "sarmd", "economic freedom", "world development report", "wdr"])
        )
        
        if is_dhs:
            consolidated[variant] = dhs_target
        elif is_lsms:
            consolidated[variant] = lsms_target
        elif is_wdi:
            consolidated[variant] = wdi_target
            
    return consolidated


# ──────────────────────────────────────────────────────────────────────────────
# Main
# ──────────────────────────────────────────────────────────────────────────────
def main():
    print("\nStep 1: Collecting raw mentions from standardized JSONs...")
    raw_df = collect_raw_mentions(STANDARDIZED_BASE, named_only=NAMED_ONLY)

    if raw_df.empty:
        print("ERROR: No mentions found. Check STANDARDIZED_BASE path and named_only filter.")
        sys.exit(1)

    print(f"\nStep 2: Preprocessing {len(raw_df)} unique raw mentions...")
    preprocessed_df = preprocess_df(raw_df)

    print(f"\nStep 3: Running clustering and hierarchization...")
    families, family_keys = run_harmonization(preprocessed_df)

    print("\nStep 3.5: Consolidating acronym families...")
    family_keys = consolidate_acronym_families(family_keys, preprocessed_df)

    print("\nStep 4: Saving outputs...")
    save_outputs(families, family_keys)

    # Print a sample of the canonical mappings for verification
    print("\nSample canonical mappings:")
    sample = list(family_keys.items())[:15]
    for variant, canonical in sample:
        print(f"  {variant!r:50s} -> {canonical!r}")

    print("\nHarmonization complete.")


if __name__ == "__main__":
    main()