"""Independent analysis of INEP Censo Escolar 2025 v2. No survey weights.

Run: python analyze.py --zip source/inep-2025.zip --out downloads
Only aggregate derived tables are emitted. Requires pandas.
"""
from pathlib import Path
import argparse, hashlib, json, zipfile
from datetime import datetime, timezone
import pandas as pd

SCHOOL = ['CO_ENTIDADE','SG_UF','NO_UF','CO_UF','CO_REGIAO','NO_REGIAO','TP_DEPENDENCIA','TP_LOCALIZACAO','TP_SITUACAO_FUNCIONAMENTO','IN_ESCOLARIZACAO','IN_BIBLIOTECA','IN_SALA_LEITURA','IN_BIBLIOTECA_SALA_LEITURA']
MAT = ['CO_ENTIDADE','QT_MAT_BAS','QT_MAT_INF','QT_MAT_FUND_AI','QT_MAT_FUND_AF','QT_MAT_MED','QT_MAT_EJA','QT_MAT_PROF']

def metric(d, mat='QT_MAT_BAS'):
    yes = d.IN_BIBLIOTECA_SALA_LEITURA.eq(1)
    no = d.IN_BIBLIOTECA_SALA_LEITURA.eq(0)
    unk = ~(yes|no)
    ns=len(d); nm=int(d[mat].sum())
    return dict(schools=ns, schools_without=int(no.sum()),schools_with=int(yes.sum()),schools_unknown=int(unk.sum()),
        enrollments=nm,enrollments_without=int(d.loc[no,mat].sum()),enrollments_with=int(d.loc[yes,mat].sum()),enrollments_unknown=int(d.loc[unk,mat].sum()),
        pct_schools_without=100*int(no.sum())/ns if ns else None,
        pct_enrollments_without=100*int(d.loc[no,mat].sum())/nm if nm else None,
        gap_school_minus_enrollment_pp=100*(int(no.sum())/ns-int(d.loc[no,mat].sum())/nm) if ns and nm else None)

def main():
    pa=argparse.ArgumentParser();pa.add_argument('--zip',default='source/inep-2025.zip');pa.add_argument('--out',default='downloads');args=pa.parse_args()
    src=Path(args.zip);out=Path(args.out);out.mkdir(parents=True,exist_ok=True)
    z=zipfile.ZipFile(src)
    def read(suffix,cols):
        n=next(x for x in z.namelist() if x.endswith(suffix))
        with z.open(n) as f: d=pd.read_csv(f,sep=';',encoding='cp1252',usecols=cols)
        assert d.CO_ENTIDADE.is_unique
        return d,n
    s,sn=read('Tabela_Escola_2025_V2.csv',SCHOOL); m,mn=read('Tabela_Matricula_2025_V2.csv',MAT)
    official_md5=z.read(next(n for n in z.namelist() if n.endswith('md5_microdados_ed_basica_2025.txt'))).decode('utf-8').splitlines()
    expected={line.split()[1].lstrip('*').lower():line.split()[0].upper() for line in official_md5 if len(line.split())==2}
    integrity=[]
    for member in [sn,mn]:
        md5=hashlib.md5(); sha=hashlib.sha256()
        with z.open(member) as f:
            while chunk:=f.read(1048576):md5.update(chunk);sha.update(chunk)
        target=expected[member.rsplit('/',1)[-1].lower()]
        assert md5.hexdigest().upper()==target, 'Source member differs from official INEP integrity manifest'
        integrity.append({'member':member,'md5':md5.hexdigest().upper(),'expected_inep_md5':target,'sha256':sha.hexdigest(),'match':True})
    repaired={}
    for key,label in [('CO_UF','SG_UF'),('CO_UF','NO_UF'),('CO_REGIAO','NO_REGIAO')]:
        lookup=s.dropna(subset=[key,label]).groupby(key)[label].agg(lambda x: x.mode().iloc[0])
        absent=s[label].isna() & s[key].notna()
        repaired[label]=int(absent.sum())
        s.loc[absent,label]=s.loc[absent,key].map(lookup)
    d=s.merge(m,on='CO_ENTIDADE',how='left',validate='one_to_one',indicator=True)
    active=d.TP_SITUACAO_FUNCIONAMENTO.eq(1)
    # School universe: active and positive basic-education enrollment. Schools
    # exclusively providing supplementary/AEE activities do not enter denominator.
    keep=active & d.QT_MAT_BAS.gt(0)
    a=d.loc[keep].copy()
    assert a[MAT[1:]].notna().all().all()
    assert (a[MAT[1:]]>=0).all().all()
    assert set(a.IN_BIBLIOTECA_SALA_LEITURA.dropna().unique()) <= {0,1}
    combined=(a.IN_BIBLIOTECA.eq(1)|a.IN_SALA_LEITURA.eq(1)).astype(int)
    observed=a.IN_BIBLIOTECA_SALA_LEITURA.notna() & a.IN_BIBLIOTECA.notna() & a.IN_SALA_LEITURA.notna()
    mismatch=int((combined[observed]!=a.loc[observed,'IN_BIBLIOTECA_SALA_LEITURA']).sum())
    assert mismatch==0
    a['network']=a.TP_DEPENDENCIA.map({1:'Federal',2:'Estadual',3:'Municipal',4:'Privada'})
    a['location']=a.TP_LOCALIZACAO.map({1:'Urbana',2:'Rural'})
    a['sector']=a.TP_DEPENDENCIA.map({1:'Pública',2:'Pública',3:'Pública',4:'Privada'})
    rows=[]
    for level,cols in [('Brasil',[]),('UF',['SG_UF']),('Região',['NO_REGIAO']),('Rede',['network']),('Localização',['location']),('Setor',['sector']),('Setor_localização',['sector','location']),('UF_pública',['SG_UF','sector'])]:
        groups=[('Brasil',a)] if not cols else a.groupby(cols,dropna=False,sort=True)
        for key,g in groups:
            if level=='UF_pública' and str(key[1])!='Pública':continue
            key=' / '.join(str(x) for x in key) if isinstance(key,tuple) else str(key)
            rows.append({'level':level,'group':key,**metric(g)})
    results=pd.DataFrame(rows);results.to_csv(out/'results.csv',index=False,float_format='%.8f')
    stage=[]
    for var,label in [('QT_MAT_INF','Educação infantil'),('QT_MAT_FUND_AI','Fundamental — anos iniciais'),('QT_MAT_FUND_AF','Fundamental — anos finais'),('QT_MAT_MED','Ensino médio'),('QT_MAT_EJA','EJA'),('QT_MAT_PROF','Educação profissional')]:
        sub=a.loc[a[var].gt(0)]
        stage.append({'stage':label,'variable':var,**metric(sub,var)})
    pd.DataFrame(stage).to_csv(out/'stages.csv',index=False,float_format='%.8f')
    # Aggregate municipality evidence, without identifying individual schools.
    # UF/state tables are the declared editorial universe; local extension later.
    audits={'download_bytes':src.stat().st_size,'download_sha256':hashlib.sha256(src.read_bytes()).hexdigest(),
       'school_file':sn,'enrollment_file':mn,'school_rows':len(s),'enrollment_rows':len(m),
       'join_school_only':int(d._merge.eq('left_only').sum()),'active_rows':int(active.sum()),
       'active_positive_enrollment_rows':len(a),'active_no_positive_enrollment':int((active&~d.QT_MAT_BAS.gt(0)).sum()),
       'total_enrollments_all_m_table':int(m.QT_MAT_BAS.sum()),'included_enrollments':int(a.QT_MAT_BAS.sum()),
       'included_unknown_library_rows':int(a.IN_BIBLIOTECA_SALA_LEITURA.isna().sum()),
       'combined_flag_mismatches':mismatch,'in_escolarizacao_distribution':a.IN_ESCOLARIZACAO.value_counts(dropna=False).to_dict(),
       'negative_count':int((a[MAT[1:]]<0).sum().sum()),'territory_display_labels_reconstructed_from_source_numeric_codes':repaired}
    samples=a.sort_values('CO_ENTIDADE').groupby(['sector','location','IN_BIBLIOTECA_SALA_LEITURA'],dropna=False).head(2)
    samples[SCHOOL+MAT[1:]].to_csv(out.parent/'sample-check.csv',index=False)
    methodology={'study_id':'brazil-school-reading-spaces-2025','created_at_utc':datetime.now(timezone.utc).isoformat(),
       'source':'INEP Censo Escolar da Educação Básica 2025 v2; updated July2026',
       'source_url':'https://www.gov.br/inep/pt-br/acesso-a-informacao/dados-abertos/microdados/censo-escolar',
       'download_url':'https://download.inep.gov.br/dados_abertos/microdados_censo_escolar_2025_.zip',
       'join':'CO_ENTIDADE one-to-one; school and enrollment aggregates',
       'universe':'TP_SITUACAO_FUNCIONAMENTO=1 and QT_MAT_BAS>0',
       'absence':'IN_BIBLIOTECA_SALA_LEITURA=0; unknown is not absence',
       'school_rate':'100 * number of included schools reporting no facility / all included schools',
       'enrollment_rate':'100 * QT_MAT_BAS at included schools reporting no facility / all included QT_MAT_BAS',
       'gap':'school rate minus enrollment rate, percentage points',
       'stage_note':'stage totals may overlap; never sum stage rows. Stage school denominator includes only schools with positive enrollment in that stage.',
       'enrollment_note':'Counts are enrollment aggregates, not an independently deduplicated count of people; QT_MAT_BAS excludes exclusive supplementary activity and exclusive AEE classes.',
       'facility_note':'Reported physical spaces existing and used at school. No collection quality, opening hours, distance, access at home or observed reading measure.',
       'reference_date':'2025-05-28 (last Wednesday of May2025)',
       'reuse':'Derived aggregate analysis and code by Alisson Bento; credit INEP for source. Source archive not redistributed; source documents retain original terms.',
       'source_member_integrity':integrity,'audits':audits}
    (out/'methodology.json').write_text(json.dumps(methodology,ensure_ascii=False,indent=2)+'\n',encoding='utf8')
    (out/'results.json').write_text(json.dumps({'audit':audits,'groups':results.to_dict(orient='records'),'stages':stage},ensure_ascii=False,indent=2)+'\n',encoding='utf8')
    print(json.dumps({'audit':audits,'main':results.loc[results.level.isin(['Brasil','Setor_localização','Rede'])].to_dict(orient='records')},ensure_ascii=False,indent=2))

if __name__=='__main__':main()