import csv,json,zipfile,io,unicodedata,itertools
from pathlib import Path
from collections import defaultdict
from fractions import Fraction
root=Path('data/clics4')
def rows(name):
 if name=='languages':return csv.DictReader(open(root/'languages.csv'))
 z=zipfile.ZipFile(root/(name+'.csv.zip'));return csv.DictReader(io.TextIOWrapper(z.open(name+'.csv'),encoding='utf-8'))
phys=['HEAVY','LIGHT (WEIGHT)','LOW','DOWN','COLD','HOT','DARK','BITTER','SWEET','HARD','SOFT','THICK','DEEP','WEAK','SLOW']
target=['SAD','GRIEF','DIFFICULT','TIRED'];want=set(phys+target)
concept={r['ID']:r['Concepticon_Gloss'] for r in rows('concepts') if r['Concepticon_Gloss'] in want}
fam={r['ID']:r['Family'] or 'isolate:'+r['ID'] for r in rows('languages')}
lex=defaultdict(lambda:defaultdict(set))
for r in rows('forms'):
 g=concept.get(r['Parameter_ID']);s=unicodedata.normalize('NFC',r['Form']).casefold().strip()
 if g and s:lex[r['Language_ID']][g].add(s)
out={}
for t in target:
 eligible=defaultdict(int);hits=set()
 for v,c in lex.items():
  if 'HEAVY' in c and t in c:
   eligible[fam[v]]+=1
   if not c['HEAVY'].isdisjoint(c[t]):hits.add(fam[v])
 out[t]={'families':len(eligible),'hits':len(hits),'max_varieties_per_family':max(eligible.values()),'hit_families':sorted(hits)}
rates=[]
for p in phys:
 den=set();num=set()
 for v,c in lex.items():
  if p in c and any(a in c for a in ['SAD','GRIEF']):
   den.add(fam[v])
   if any(not c[p].isdisjoint(c[a]) for a in ['SAD','GRIEF'] if a in c):num.add(fam[v])
 rates.append(Fraction(len(num),len(den)))
obs=sum(rates[phys.index(p)] for p in ['HEAVY','LOW','DOWN'])/3-(sum(rates)-sum(rates[phys.index(p)] for p in ['HEAVY','LOW','DOWN']))/12
vals=[sum(rates[i] for i in ids)/3-(sum(rates)-sum(rates[i] for i in ids))/12 for ids in itertools.combinations(range(15),3)]
out['exact_test']={'assignments':len(vals),'at_least_observed':sum(x>=obs for x in vals),'p':sum(x>=obs for x in vals)/len(vals),'difference_unrounded':float(obs)}
Path('results').mkdir(exist_ok=True);Path('results/independent.json').write_text(json.dumps(out,indent=2));print(json.dumps(out,indent=2))