district_mapper.py

"""
district_mapper.py
──────────────────
District back-mapping engine for KYS school scraper.

Two responsibilities:
  1. build_mapper(analysis_dict)       – builds lookup tables from the analysis results
  2. apply_geo_decode(df, geo_lookup)  – fills real State/District/Block for IAF/KVS/NVS/NAVY rows
  3. apply_district_backmap(df, mapper)– maps new district names → old district names

Non-actual states that need geo decoding:
  IAF EC SOCIETY, KENDRIYA VIDYALAYA SANGHATHAN,
  NAVODAYA VIDYALAYA SAMITI, NAVY EDUCATION SOCIETY, MSRVVP
"""

import pandas as pd
from collections import defaultdict

# ── Non-actual state names that need UDISE geo-decoding ──────────────────────
NON_ACTUAL_STATES = {
    "IAF EC SOCIETY",
    "KENDRIYA VIDYALAYA SANGHATHAN",
    "NAVODAYA VIDYALAYA SAMITI",
    "NAVY EDUCATION SOCIETY",
    "MSRVVP",
}

# ── UDISE state-code → state name (standard DISE/UDISE coding) ───────────────
UDISE_STATE_CODES = {
    "01": "JAMMU & KASHMIR",
    "02": "HIMACHAL PRADESH",
    "03": "PUNJAB",
    "04": "CHANDIGARH",
    "05": "UTTARAKHAND",
    "06": "HARYANA",
    "07": "DELHI",
    "08": "RAJASTHAN",
    "09": "UTTAR PRADESH",
    "10": "BIHAR",
    "11": "SIKKIM",
    "12": "ARUNACHAL PRADESH",
    "13": "NAGALAND",
    "14": "MANIPUR",
    "15": "MIZORAM",
    "16": "TRIPURA",
    "17": "MEGHALAYA",
    "18": "ASSAM",
    "19": "WEST BENGAL",
    "20": "JHARKHAND",
    "21": "ODISHA",
    "22": "CHHATTISGARH",
    "23": "MADHYA PRADESH",
    "24": "GUJARAT",
    "25": "DADRA & NAGAR HAVELI AND DAMAN & DIU",
    "26": "DADRA & NAGAR HAVELI AND DAMAN & DIU",
    "27": "MAHARASHTRA",
    "28": "ANDHRA PRADESH",
    "29": "KARNATAKA",
    "30": "GOA",
    "31": "LAKSHADWEEP",
    "32": "KERALA",
    "33": "TAMILNADU",
    "34": "PUDUCHERRY",
    "35": "ANDAMAN & NICOBAR ISLANDS",
    "36": "TELANGANA",
    "37": "LADAKH",
}


# ─────────────────────────────────────────────────────────────────────────────
# build_mapper
# ─────────────────────────────────────────────────────────────────────────────
def build_mapper(analysis_dict: dict) -> dict:
    """
    Build lookup tables from the analysis results dict.

    Parameters
    ----------
    analysis_dict : dict
        Keys: "District Mapping", "Complex Splits" (DataFrames from run_analysis or Excel)

    Returns
    -------
    mapper : dict with keys:
        simple_map   – { (state, new_dist) -> old_dist }       for 1-to-1 renames
        complex_map  – { udise_str -> old_dist }               per-school UDISE lookup
        block_map    – { (state, new_dist, new_block) -> old_dist }  block-level fallback
        dominant_map – { (state, new_dist) -> most_common_old_dist } last-resort fallback
    """
    dm = analysis_dict.get("District Mapping", pd.DataFrame())
    b2d = analysis_dict.get("Block to District", pd.DataFrame())
    bm = analysis_dict.get("Block Mapping", pd.DataFrame())

    simple_map   = {}
    dominant_map = {}
    complex_map  = {}
    block_map    = {}
    block_rename = {}

    # ── 1. Simple map: (state, new_dist) → old_dist ──────────────────────────
    if not dm.empty:
        for _, row in dm.iterrows():
            state    = str(row.get("State", "")).strip().upper()
            scraped_dist = str(row.get("Dist_OLD", "")).strip().upper()
            sf_dist = str(row.get("Dist_NEW", "")).strip().upper()
            if not state or not scraped_dist or not sf_dist:
                continue
            key = (state, scraped_dist)
            if key not in simple_map:
                simple_map[key] = sf_dist
            # Track dominant (most schools) for ambiguous cases
            val = row.get("Affected_Schools", 0)
            count = int(val) if pd.notna(val) else 0
            if key not in dominant_map or count > dominant_map[key][1]:
                dominant_map[key] = (sf_dist, count)
                
            # Per-UDISE lookup (highest accuracy) - populate for ALL common schools
            udise_raw = row.get("List_of_UDISEs", "")
            if pd.notna(udise_raw) and str(udise_raw).strip():
                for u in str(udise_raw).split(","):
                    u = u.strip().replace(".0", "").zfill(11)
                    if u:
                        complex_map[u] = sf_dist

    # Strip count from dominant_map
    dominant_map = {k: v[0] for k, v in dominant_map.items()}

    # ── 2. Block-level fallback mapping ──────────────────────────────────────
    if not b2d.empty:
        for _, row in b2d.iterrows():
            state    = str(row.get("State", "")).strip().upper()
            old_dist = str(row.get("Dist_OLD", "")).strip().upper()
            new_dist = str(row.get("Dist_NEW", "")).strip().upper()
            new_block= str(row.get("Block_NEW", "")).strip().upper()

            # Block-level fallback
            bkey = (state, new_dist, new_block)
            if bkey not in block_map: # block_to_district is sorted by Schools DESC, so first is dominant
                block_map[bkey] = old_dist

    # ── 3. Block renames ─────────────────────────────────────────────────────────
    if not bm.empty:
        for _, row in bm.iterrows():
            state     = str(row.get("State", "")).strip().upper()
            dist_old  = str(row.get("Dist_OLD", "")).strip().upper()
            block_old = str(row.get("Block_OLD", "")).strip().upper()
            block_new = str(row.get("Block_NEW", "")).strip().upper()
            
            # Map by (state, new_block) to old_block. 
            # We also include dist_old in the key for safety since blocks can share names.
            if state and dist_old and block_new and block_old:
                block_rename[(state, dist_old, block_new)] = block_old

    return {
        "simple_map":   simple_map,
        "dominant_map": dominant_map,
        "complex_map":  complex_map,
        "block_map":    block_map,
        "block_rename": block_rename,
    }


# ─────────────────────────────────────────────────────────────────────────────
# build_udise_geo_lookup
# ─────────────────────────────────────────────────────────────────────────────
def build_udise_geo_lookup(master_df: pd.DataFrame) -> dict:
    """
    Build UDISE prefix → real State/District/Block lookup from master data.

    Uses only rows from REAL states (not IAF/KVS/NVS/NAVY) so the lookup is
    grounded in actual geographic data.

    Returns
    -------
    dict with keys:
        "state"    – { "SS"     → state_name }
        "district" – { "SSDD"   → district_name }
        "block"    – { "SSDDBB" → block_name }
    """
    geo = {"state": {}, "district": {}, "block": {}}

    if master_df is None or master_df.empty:
        # Fall back to hard-coded UDISE state codes only
        geo["state"] = UDISE_STATE_CODES.copy()
        return geo

    # Use only real-state rows for building the geo lookup
    real_rows = master_df[
        ~master_df["School_State__c"].str.strip().str.upper().isin(NON_ACTUAL_STATES)
    ].copy()

    udise_col = None
    for c in ["School_Udise_Code__c", "UDISE"]:
        if c in real_rows.columns:
            udise_col = c
            break

    if udise_col is None:
        geo["state"] = UDISE_STATE_CODES.copy()
        return geo

    real_rows["_udise"] = (
        real_rows[udise_col]
        .astype(str)
        .str.replace(r"\.0$", "", regex=True)
        .str.zfill(11)
    )
    real_rows["_state"]    = real_rows["School_State__c"].str.strip().str.upper()
    real_rows["_district"] = real_rows["School_District__c"].str.strip().str.upper()
    real_rows["_block"]    = real_rows["School_Block__c"].str.strip().str.upper()

    for _, row in real_rows.iterrows():
        u = row["_udise"]
        if len(u) < 6:
            continue
        ss    = u[:2]
        ssdd  = u[:4]
        ssddbb= u[:6]

        if ss not in geo["state"]    and row["_state"]:
            geo["state"][ss]       = row["_state"]
        if ssdd not in geo["district"] and row["_district"]:
            geo["district"][ssdd]  = row["_district"]
        if ssddbb not in geo["block"]  and row["_block"]:
            geo["block"][ssddbb]   = row["_block"]

    # Ensure base UDISE state codes are always present
    for code, name in UDISE_STATE_CODES.items():
        geo["state"].setdefault(code, name)

    return geo


# ─────────────────────────────────────────────────────────────────────────────
# apply_geo_decode
# ─────────────────────────────────────────────────────────────────────────────
def apply_geo_decode(df: pd.DataFrame, geo_lookup: dict) -> tuple:
    """
    For rows belonging to non-actual states (IAF, KVS, NVS, NAVY, MSRVVP),
    decode the real State / District / Block from the UDISE code prefix and
    overwrite those columns.

    Returns (mapped_df, report_df)
    """
    df = df.copy()

    udise_col = None
    for c in ["School_Udise_Code__c", "UDISE"]:
        if c in df.columns:
            udise_col = c
            break
    if udise_col is None:
        return df, pd.DataFrame()  # nothing to do

    mask = df["School_State__c"].str.strip().str.upper().isin(NON_ACTUAL_STATES)
    if not mask.any():
        return df, pd.DataFrame()

    udise_series = (
        df.loc[mask, udise_col]
        .astype(str)
        .str.replace(r"\.0$", "", regex=True)
        .str.zfill(11)
    )

    changes = []
    
    # Iterate over the masked rows to track changes
    for idx, orig_state in df.loc[mask, "School_State__c"].items():
        u = udise_series[idx]
        orig_dist = df.at[idx, "School_District__c"]
        
        new_state = geo_lookup.get("state", {}).get(u[:2], "") if len(u) >= 2 else ""
        new_state = new_state if new_state else orig_state
        
        new_dist = geo_lookup.get("district", {}).get(u[:4], "") if len(u) >= 4 else ""
        new_dist = new_dist if new_dist else orig_dist
        
        new_block = geo_lookup.get("block", {}).get(u[:6], "") if len(u) >= 6 else ""
        new_block = new_block if new_block else df.at[idx, "School_Block__c"]
        
        if new_dist != orig_dist:
            changes.append({
                "State": orig_state, # The original virtual state (e.g. NAVY)
                "Type": "District (UDISE Decode)",
                "New_Value": new_dist,
                "Old_Value": orig_dist
            })
            
        df.at[idx, "School_State__c"] = new_state
        df.at[idx, "School_District__c"] = new_dist
        df.at[idx, "School_Block__c"] = new_block

    return df, pd.DataFrame(changes)


# ─────────────────────────────────────────────────────────────────────────────
# apply_district_backmap
# ─────────────────────────────────────────────────────────────────────────────
def apply_district_backmap(df: pd.DataFrame, mapper: dict) -> tuple:
    """
    Map every row's School_District__c from new name → old name.

    Resolution order for districts:
      1. Per-school UDISE lookup (complex_map)        — highest accuracy
      2. Block-level lookup (block_map)               — for new schools in split districts
      3. Simple 1:1 rename (simple_map)               — for ordinary renames
      4. Dominant old district (dominant_map)         — last resort
      
    Resolution for blocks:
      - Uses block_rename map matching on (state, mapped_district, new_block)

    Returns
    -------
    (mapped_df, report_df)
      mapped_df  – DataFrame with mapped District and Block
      report_df  – Summary of what changed
    """
    df = df.copy()

    udise_col = None
    for c in ["School_Udise_Code__c", "UDISE"]:
        if c in df.columns:
            udise_col = c
            break

    simple_map   = mapper.get("simple_map",   {})
    complex_map  = mapper.get("complex_map",  {})
    block_map    = mapper.get("block_map",    {})
    dominant_map = mapper.get("dominant_map", {})
    block_rename = mapper.get("block_rename", {})

    changes = []

    for idx, row in df.iterrows():
        state    = str(row.get("School_State__c", "")).strip().upper()
        new_dist = str(row.get("School_District__c", "")).strip().upper()
        new_block= str(row.get("School_Block__c", "")).strip().upper()

        old_dist = None

        # 1. UDISE per-school lookup
        if udise_col:
            u = (
                str(row.get(udise_col, ""))
                .replace(".0", "")
                .zfill(11)
            )
            old_dist = complex_map.get(u)

        # 2. Block-level fallback
        if old_dist is None:
            old_dist = block_map.get((state, new_dist, new_block))

        # 3. Simple 1:1 rename
        if old_dist is None:
            old_dist = simple_map.get((state, new_dist))

        # 4. Dominant fallback
        if old_dist is None:
            old_dist = dominant_map.get((state, new_dist))

        if old_dist and old_dist != new_dist:
            changes.append({
                "State":        state,
                "Type":         "District",
                "New_Value":    new_dist,
                "Old_Value":    old_dist,
            })
            df.at[idx, "School_District__c"] = old_dist
            
        final_dist = old_dist if old_dist else new_dist
        
        # 5. Block rename
        # Disabled as per user request: "only district names changed not blocks ntg else"
        # old_block = block_rename.get((state, final_dist, new_block))
        # if old_block and old_block != new_block:
        #     changes.append({
        #         "State":        state,
        #         "Type":         "Block",
        #         "New_Value":    new_block,
        #         "Old_Value":    old_block,
        #     })
        #     df.at[idx, "School_Block__c"] = old_block

    # Build change report
    if changes:
        report_df = (
            pd.DataFrame(changes)
            .groupby(["State", "Type", "New_Value", "Old_Value"])
            .size()
            .reset_index(name="Schools_Remapped")
            .sort_values(["State", "Type", "Schools_Remapped"], ascending=[True, True, False])
        )
    else:
        report_df = pd.DataFrame(
            columns=["State", "Type", "New_Value", "Old_Value", "Schools_Remapped"]
        )

    return df, report_df