#!/usr/bin/env python3
"""Movement Thinking 3.0: design calculations and explicitly scoped data audits.

No command recruits participants, transmits records, joins TALIS to SEDA, or
claims a completed field study. All input/output paths are investigator supplied.
"""
import argparse
import csv
import hashlib
import itertools
import json
import math
from pathlib import Path

import numpy as np
from scipy.stats import norm, t

VERSION = '3.0'


def write(path, value):
    Path(path).write_text(json.dumps(value, indent=2, allow_nan=False) + '\n')


def design(replicates=10000, seed=20261015):
    if not 100 <= replicates <= 1000000:
        raise ValueError('Use 100 to 1,000,000 simulation replicates.')
    rng = np.random.default_rng(seed)
    rows = []
    for schools, rho in itertools.product([48, 80, 120], [.05, .10, .15]):
        blocks, size = schools // 4, 40
        z = np.argsort(np.argsort(rng.random((replicates, blocks, 4)), axis=2), axis=2) < 2
        errors = rng.normal(0, math.sqrt(rho + (1-rho)/size), (replicates, blocks, 4))
        critical = t.ppf(.975, blocks-1)
        rejection = {}
        for delta in [0., .25]:
            y = errors + delta*z
            yt = y[z].reshape(replicates, blocks, 2)
            yc = y[~z].reshape(replicates, blocks, 2)
            effect = (yt.mean(2)-yc.mean(2)).mean(1)
            variance = ((yt.var(2, ddof=1)/2 + yc.var(2, ddof=1)/2).sum(1))/blocks**2
            rate = float(np.mean(np.abs(effect/np.sqrt(variance)) > critical))
            rejection[str(delta)] = {'rejection_rate':rate, 'monte_carlo_se':math.sqrt(rate*(1-rate)/replicates)}
        de = 1 + (size-1)*rho
        rows.append({'schools':schools,'district_blocks':blocks,'students_per_school':size,
                     'planned_students':schools*size,'icc_of_change_score':rho,'design_effect':de,
                     'normal_approximate_mde_sd_change':float((norm.ppf(.975)+norm.ppf(.80))*math.sqrt(4*de/(schools*size))),
                     'simulation':rejection})
    return {'version':VERSION,'data_kind':'simulation_only','seed':seed,'replicates':replicates,
            'assumptions':['Four schools per block; two assigned to each arm.',
                           'Independent Gaussian school mean changes; individual change-score SD=1.',
                           'Constant additive effect, equal cluster sizes, complete outcomes, no interference.',
                           'Neyman within-block variance; two-sided t critical value with blocks-1 df.'],
            'not_estimated':['Movement Thinking treatment effects','actual ICCs','actual school or student outcomes'],
            'scenarios':rows}


def audit_seda(path):
    data = Path(path).read_bytes()
    with Path(path).open(encoding='utf-8-sig', newline='') as stream:
        reader = csv.DictReader(stream)
        required = {'sedaadmin','subcat','subgroup','gap','multi_comp','year_min','year_max','pooled_years',
                    'cs_mn_lrn_ol','cs_mn_lrn_ol_se_adj'}
        if not required <= set(reader.fieldnames or []):
            raise ValueError('This adapter requires the documented SEDA 2025.2 administrative-district pooled CS schema.')
        rows = list(reader)
    keys = ['sedaadmin','subcat','subgroup','gap','multi_comp','year_min','year_max']
    unique = len({tuple(r[k] for k in keys) for r in rows})
    if unique != len(rows):
        raise ValueError('Duplicate full release keys: investigate before analysis.')
    cohort = [r for r in rows if r['subcat']=='all' and r['subgroup']=='all' and r['gap']=='0'
              and r['multi_comp']=='0' and r['pooled_years']=='2022-2025']
    if len({r['sedaadmin'] for r in cohort}) != len(cohort):
        raise ValueError('The selected context slice is not one record per district.')
    def observed(r):
        try:
            v, se = float(r['cs_mn_lrn_ol']), float(r['cs_mn_lrn_ol_se_adj'])
            return math.isfinite(v) and math.isfinite(se) and se > 0
        except ValueError:
            return False
    return {'data_kind':'secondary_data_feasibility_audit','version':VERSION,
            'dataset':'SEDA 2025.2; seda_admindist_pool_cs_2025.2.csv',
            'source_url':'https://stacks.stanford.edu/file/np279jm6134/seda_admindist_pool_cs_2025.2.csv',
            'sha256':hashlib.sha256(data).hexdigest(),'checked_on':'2026-10-15',
            'rows':len(rows),'distinct_administrative_district_ids':len({r['sedaadmin'] for r in rows}),
            'distinct_admin_subgroup_pairs':len({(r['sedaadmin'],r['subgroup']) for r in rows}),
            'full_key_fields':keys,'unique_full_keys':unique,
            'context_filter':{'subcat':'all','subgroup':'all','gap':'0','multi_comp':'0','pooled_years':'2022-2025'},
            'selected_context_rows':len(cohort),'selected_with_rate_and_positive_adjusted_se':sum(map(observed,cohort)),
            'selected_missing_rate_or_valid_adjusted_se':sum(not observed(r) for r in cohort),
            'interpretation':'File/schema feasibility only. No TALIS linkage, local participant data, intervention effect or causal validation.'}


def randomize(roster, seed):
    schools = roster['schools']
    if roster.get('data_kind') not in ['synthetic_demonstration','approved_study_data']:
        raise ValueError('Declare synthetic or governed study-data provenance.')
    if roster['data_kind']=='approved_study_data' and len(schools)!=80:
        raise ValueError('Version 3.0 specifies 80 schools. A different actual study design requires a protocol amendment.')
    if len({s['school_id'] for s in schools}) != len(schools):
        raise ValueError('School IDs must be globally unique.')
    rng = np.random.default_rng(seed)
    out = []
    for district in sorted({s['district_id'] for s in schools}):
        block = sorted([s for s in schools if s['district_id']==district], key=lambda s:s['school_id'])
        if len(block) != 4:
            raise ValueError('The registered allocation requires exactly four schools per district block.')
        treated = set(rng.choice(4, size=2, replace=False).tolist())
        out.extend({**s,'assignment':int(i in treated)} for i,s in enumerate(block))
    return {'data_kind':roster['data_kind'],'protocol_version':VERSION,'seed':seed,'schools':out}


def analyze_complete_cohort(data, permutations=9999, seed=20261015):
    if not 99 <= permutations <= 1000000:
        raise ValueError('Use 99 to 1,000,000 Monte Carlo allocations.')
    if data.get('data_kind') not in ['synthetic_demonstration','approved_study_data']:
        raise ValueError('Label the provenance of the data explicitly.')
    if data.get('protocol_version') != VERSION:
        raise ValueError('Protocol version mismatch.')
    schools, rows = data['schools'], data['students']
    if data['data_kind']=='approved_study_data' and len(schools)!=80:
        raise ValueError('Actual study data do not match the 80-school version 3.0 design.')
    ids = [s['school_id'] for s in schools]
    if len(ids)!=len(set(ids)) or len({r['student_id'] for r in rows})!=len(rows):
        raise ValueError('Duplicate school or student IDs.')
    if any(r['school_id'] not in ids for r in rows):
        raise ValueError('A student record references an unknown school.')
    means = {}
    for school in schools:
        own = [r for r in rows if r['school_id']==school['school_id']]
        if len(own) != 40:
            raise ValueError('Exactly 40 baseline-cohort students per school are required by this complete-cohort adapter.')
        changes = []
        for row in own:
            values = [row.get('baseline_score'),row.get('endline_score')]
            if any(isinstance(v,bool) or not isinstance(v,(float,int)) or not math.isfinite(v) or not 0<=v<=20 for v in values):
                raise ValueError('Missing or invalid scores: use the prespecified missing-data process; do not delete records to run this adapter.')
            changes.append(values[1]-values[0])
        means[school['school_id']] = float(np.mean(changes))
    y,z = [],[]
    for district in sorted({s['district_id'] for s in schools}):
        own = sorted([s for s in schools if s['district_id']==district],key=lambda s:s['school_id'])
        a = [s['assignment'] for s in own]
        if len(own)!=4 or sorted(a)!=[0,0,1,1]:
            raise ValueError('Expected two treated and two controls in every four-school block.')
        y.append([means[s['school_id']] for s in own]);z.append(a)
    y,z = np.asarray(y),np.asarray(z,dtype=bool)
    blocks = len(y)
    if blocks<2:
        raise ValueError('At least two district blocks are needed.')
    yt,yc = y[z].reshape(blocks,2),y[~z].reshape(blocks,2)
    effect = float(np.mean(yt.mean(1)-yc.mean(1)))
    se = float(np.sqrt(np.sum(yt.var(1,ddof=1)/2 + yc.var(1,ddof=1)/2)/blocks**2))
    if se<=0:
        raise ValueError('Degenerate between-school variation: obtain statistical review instead of a zero-width uncertainty interval.')
    critical = float(t.ppf(.975,blocks-1))
    rng = np.random.default_rng(seed)
    samples = np.argsort(np.argsort(rng.random((permutations,blocks,4)),axis=2),axis=2)<2
    sampled_effects = np.mean(np.sum(y*samples,axis=2)/2 - np.sum(y*(~samples),axis=2)/2,axis=1)
    p = float((1+np.count_nonzero(np.abs(sampled_effects)>=abs(effect)-1e-12))/(permutations+1))
    return {'data_kind':data['data_kind'],'protocol_version':VERSION,'schools':len(schools),'students':len(rows),
            'estimand':'Equal-school average baseline-cohort change-score ITT effect, rubric points.',
            'effect_points':effect,'neyman_se':se,'approximate_95_ci':[effect-critical*se,effect+critical*se],
            'ci_df':blocks-1,'sharp_null_randomization_p':p,'permutations':permutations,'seed':seed,
            'limits':'CI uses a conservative small-sample t approximation; randomization p tests the sharp no-effect null. This adapter requires complete cohorts and does not implement missing-data imputation or fit a mediation model.'}


def talis_brr_proportion(path, country, variable, category):
    if variable not in ['TT4G26A','TT4G26B','TT4G26C','TT4G26D','TT4G26E','TT4G26F','TT4G26G','TT4G26H'] or category not in range(1,7):
        raise ValueError('This adapter is restricted to documented question-26 items and response categories 1-6.')
    names = ['TCHWGT']+[f'TRWGT{i}' for i in range(1,101)]
    numerator,denominator = np.zeros(101),np.zeros(101)
    n,missing = 0,0
    with Path(path).open(encoding='utf-8-sig',newline='') as stream:
        reader=csv.DictReader(stream)
        if not set(names+['CNTRY','IDPOP',variable]) <= set(reader.fieldnames or []):
            raise ValueError('The input must contain the documented teacher weight/replicate columns and population identifiers.')
        for row in reader:
            if row['CNTRY']!=country or float(row['IDPOP'])!=2:continue
            try:value=float(row[variable])
            except ValueError:missing+=1;continue
            if value not in range(1,7):missing+=1;continue
            weights=np.array([float(row[k]) for k in names])
            if not np.isfinite(weights).all() or (weights<0).any() or weights[0]<=0:
                raise ValueError('Invalid survey weights.')
            numerator+=weights*(value==category);denominator+=weights;n+=1
    if (denominator<=0).any():raise ValueError('No valid population domain or a zero replicate denominator.')
    estimates=numerator/denominator
    se=float(np.sqrt(np.sum((estimates[1:]-estimates[0])**2)/(100*(1-.5)**2)))
    return {'data_kind':'secondary_TALIS_analysis','country':country,'idpop':2,'variable':variable,'category':category,
            'valid_responses':n,'missing_or_non_substantive':missing,'weighted_proportion':float(estimates[0]),
            'fay_brr_se':se,'replicates':100,'fay_factor':.5,
            'participation_qualification':'The official guide rates U.S. teacher participation Poor.' if country=='USA' else 'Consult the country/population adjudication in the official guide.',
            'limits':'Descriptive valid-response proportion, not a causal effect, a joint TALIS-SEDA model or a validated Movement Thinking scale.'}


def main():
    p=argparse.ArgumentParser(description=__doc__);sub=p.add_subparsers(dest='command',required=True)
    s=sub.add_parser('design');s.add_argument('--replicates',type=int,default=10000);s.add_argument('--seed',type=int,default=20261015);s.add_argument('--output',required=True)
    s=sub.add_parser('audit-seda');s.add_argument('--input',required=True);s.add_argument('--output',required=True)
    s=sub.add_parser('randomize');s.add_argument('--input',required=True);s.add_argument('--seed',type=int,required=True);s.add_argument('--output',required=True)
    s=sub.add_parser('analyze-complete');s.add_argument('--input',required=True);s.add_argument('--permutations',type=int,default=9999);s.add_argument('--output',required=True)
    s=sub.add_parser('talis-brr');s.add_argument('--input',required=True);s.add_argument('--country',required=True);s.add_argument('--variable',required=True);s.add_argument('--category',type=int,required=True);s.add_argument('--output',required=True)
    a=p.parse_args()
    if a.command=='design':out=design(a.replicates,a.seed)
    elif a.command=='audit-seda':out=audit_seda(a.input)
    elif a.command=='randomize':out=randomize(json.loads(Path(a.input).read_text()),a.seed)
    elif a.command=='analyze-complete':out=analyze_complete_cohort(json.loads(Path(a.input).read_text()),a.permutations)
    else:out=talis_brr_proportion(a.input,a.country,a.variable,a.category)
    write(a.output,out)
    print(f'{a.command}: saved {out.get("data_kind", "specified analysis")} output to {a.output}')


if __name__=='__main__':main()