import pandas as pd
import os

def is_govt(management_name):
    """Simple function to classify if a school is strictly Government managed"""
    val = str(management_name).lower()
    if 'aided' in val or 'private' in val or 'unaided' in val or 'un-aided' in val:
        return False
    return 'gov' in val or 'dept of edu' in val or 'local body' in val or 'panchayat' in val or 'municipal' in val

def run_analysis():
    # Define paths (Update the parquet path to point to a real file when testing)
    csv_path = "master_all_states.csv"
    parquet_path = "mapped_master_2026_01.parquet" # Replace with your actual parquet file path

    if not os.path.exists(csv_path) or not os.path.exists(parquet_path):
        print(f"Make sure both {csv_path} and {parquet_path} exist in this folder to run the script!")
        return

    # 1. Load the 2025 dataset and the Newest dataset
    print("Loading datasets...")
    df_2025 = pd.read_csv(csv_path, low_memory=False)
    df_new = pd.read_parquet(parquet_path)

    # 2. Normalize UDISE codes to ensure perfect 11-digit matching
    df_2025['_U'] = df_2025['School_Udise_Code__c'].astype(str).str.strip().str.zfill(11)
    df_new['_U'] = df_new['School_Udise_Code__c'].astype(str).str.strip().str.zfill(11)

    # 3. Apply Government Flag using the strict School_Management_Type__c column
    df_2025['is_govt'] = df_2025['School_Management_Type__c'].apply(is_govt)
    df_new['is_govt'] = df_new['School_Management_Type__c'].apply(is_govt)

    print(f"Total Govt Schools (2025): {df_2025['is_govt'].sum():,}")
    print(f"Total Govt Schools (New):  {df_new['is_govt'].sum():,}")

    # ---------------------------------------------------------
    # FINDING MISSING SCHOOLS
    # ---------------------------------------------------------
    udises_2025 = set(df_2025['_U'].dropna())
    udises_new = set(df_new['_U'].dropna())
    missing_udises = udises_2025 - udises_new
    
    # Filter 2025 data to only show missing schools, and DROP duplicates for accurate counts!
    missing_df = df_2025[df_2025['_U'].isin(missing_udises)].drop_duplicates(subset=['_U']).copy()
    
    # Add Master Year column
    missing_df.insert(0, "Missing In Master", "2026") 
    
    print(f"\nTotal Missing Schools: {len(missing_df)}")
    print("Missing Schools by State:\n", missing_df['School_State__c'].value_counts().head())

    # ---------------------------------------------------------
    # FINDING MANAGEMENT SHIFTS (Govt -> Non-Govt)
    # ---------------------------------------------------------
    # Merge the two datasets on UDISE code
    merged = pd.merge(
        df_2025[['_U', 'School_Name__c', 'School_Management_Type__c', 'is_govt']], 
        df_new[['_U', 'School_Name__c', 'School_Management_Type__c', 'is_govt']], 
        on='_U', 
        suffixes=('_old', '_new')
    )

    # Filter where it WAS govt, but is NO LONGER govt
    changed_mgmt = merged[(merged['is_govt_old'] == True) & (merged['is_govt_new'] == False)].copy()
    changed_mgmt.insert(0, "Shifted In Master", "2026")

    print(f"\nSchools shifted from Govt -> Non-Govt: {len(changed_mgmt)}")
    return missing_df, changed_mgmt

if __name__ == "__main__":
    run_analysis()