"""CPU-only learned teacher/student classification distillation. Python >=3.9, stdlib."""
import copy
import json
import math
import platform
import random
import statistics
from pathlib import Path


def softmax(z, temperature=1.0):
    if temperature <= 0:
        raise ValueError('temperature must be positive')
    a = [v / temperature for v in z]
    m = max(a)
    e = [math.exp(v - m) for v in a]
    return [v / sum(e) for v in e]


def logits(w, x):
    return [sum(a * b for a, b in zip(row, x)) for row in w]


def argmax(p):
    return max(range(len(p)), key=p.__getitem__)


def kl(p, q):
    return sum(a * math.log(a / b) for a, b in zip(p, q) if a)


def features(x, teacher=False):
    a, b = x
    return [1.0, a, b, a * a, a * b, b * b] if teacher else [1.0, a, b]


def dataset(seed, n):
    r = random.Random(seed)
    out = []
    for _ in range(n):
        a, b = r.uniform(-2, 2), r.uniform(-2, 2)
        y = argmax([1.1*a + .9*b + .8*a*b, -a + .6*b - .6*a*b, -.9*b + .45*a*a - .4])
        out.append({'x': [a, b], 'y': y})
    return out


def init(seed, d):
    r = random.Random(seed)
    return [[r.gauss(0, .02) for _ in range(d)] for _ in range(3)]


def train(w, rows, teacher_features=False, teacher=None, mode='gold', temperature=1., alpha=.7, steps=250):
    w = copy.deepcopy(w)
    target = []
    for row in rows:
        if mode == 'gold':
            target.append([float(c == row['y']) for c in range(3)])
        else:
            z = logits(teacher, features(row['x'], True))
            target.append(softmax(z, temperature) if mode == 'soft' else [float(c == argmax(z)) for c in range(3)])
    history = []
    for step in range(steps):
        grad = [[0.] * len(w[0]) for _ in range(3)]
        loss = 0.
        for row, p in zip(rows, target):
            x = features(row['x'], teacher_features)
            z = logits(w, x)
            q1 = softmax(z)
            if mode == 'soft':
                q = softmax(z, temperature)
                loss += (1-alpha)*(-math.log(q1[row['y']])) + alpha*temperature**2*kl(p, q)
                dz = [(1-alpha)*(q1[c]-(c == row['y'])) + alpha*temperature*(q[c]-p[c]) for c in range(3)]
            else:
                loss -= sum(p[c]*math.log(q1[c]) for c in range(3))
                dz = [q1[c]-p[c] for c in range(3)]
            for c in range(3):
                for j in range(len(x)):
                    grad[c][j] += dz[c]*x[j]
        if step in (0, steps-1):
            history.append({'step_before_update': step, 'objective': loss/len(rows)})
        for c in range(3):
            for j in range(len(w[0])):
                w[c][j] -= .12*grad[c][j]/len(rows)
    return w, history


def evaluate(w, rows, teacher, is_teacher=False):
    pred = [softmax(logits(w, features(r['x'], is_teacher))) for r in rows]
    tp = [softmax(logits(teacher, features(r['x'], True))) for r in rows]
    wrong = [i for i, r in enumerate(rows) if argmax(tp[i]) != r['y']]
    return {'accuracy': statistics.mean(argmax(p) == r['y'] for p, r in zip(pred, rows)),
            'gold_nll': statistics.mean(-math.log(p[r['y']]) for p, r in zip(pred, rows)),
            'teacher_agreement': statistics.mean(argmax(p) == argmax(t) for p,t in zip(pred,tp)),
            'teacher_kl_t1': statistics.mean(kl(t,p) for t,p in zip(tp,pred)),
            'teacher_wrong_count': len(wrong),
            'copied_teacher_error_rate': (statistics.mean(argmax(pred[i]) == argmax(tp[i]) for i in wrong) if wrong else None)}


def checks():
    p = softmax([2., -.5, 1.], 3.)
    z = [.3, -.2, .9]
    q = softmax(z, 3.)
    analytic = [3*(b-a) for a,b in zip(p,q)]
    numeric=[]
    for c in range(3):
        hi,lo=z[:],z[:]
        hi[c]+=1e-5; lo[c]-=1e-5
        numeric.append((9*kl(p,softmax(hi,3.))-9*kl(p,softmax(lo,3.)))/2e-5)
    error=max(abs(a-b) for a,b in zip(analytic,numeric))
    assert error < 1e-8
    assert abs(kl(p,p)) < 1e-12
    assert abs(sum(softmax([1000.,1001.,999.]))-1) < 1e-12
    assert abs(kl([.8,.1,.1],[.4,.3,.3])-kl([.4,.3,.3],[.8,.1,.1])) > .01
    return {'finite_difference_max_error': error, 'kl_self_zero': True, 'kl_direction_distinct': True, 'large_logit_stability': True}


def main():
    data={name:dataset(seed,n) for name,seed,n in [('teacher_train',6600,360),('transfer',6601,90),('test',6602,240)]}
    assert not ({tuple(r['x']) for r in data['teacher_train']} & {tuple(r['x']) for r in data['test']})
    runs=[]
    for seed in [66,67,68]:
        teacher,th=train(init(seed,6),data['teacher_train'],teacher_features=True,steps=400)
        frozen=copy.deepcopy(teacher)
        shared=init(seed+100,3)
        variants=[('gold','gold',1.),('teacher_argmax','hard',1.),('soft_t1','soft',1.),('soft_t3','soft',3.)]
        row={'seed':seed,'teacher':evaluate(teacher,data['test'],teacher,True),'teacher_history':th,'teacher_weights':teacher,'students':{}}
        for name,mode,t in variants:
            student,h=train(shared,data['transfer'],teacher=teacher,mode=mode,temperature=t)
            row['students'][name]={'metrics':evaluate(student,data['test'],teacher),'history':h,'weights':student}
        assert teacher==frozen
        runs.append(row)
    aggregate={}
    for name in runs[0]['students']:
        aggregate[name]={}
        for metric in runs[0]['students'][name]['metrics']:
            vals=[r['students'][name]['metrics'][metric] for r in runs]
            aggregate[name][metric]={'mean':statistics.mean(vals),'sample_sd':statistics.stdev(vals)}
    result={'python':platform.python_version(),'device':'CPU','seeds':[66,67,68],
            'protocol':{'teacher_steps':400,'student_steps':250,'lr':.12,'alpha_soft':.7,'selection':'fixed final step; no hyperparameter selection','shapes':{'teacher_X':[360,6],'student_X':[90,3],'test_logits':[240,3]},'teacher_parameters':18,'student_parameters':9},
            'checks':checks(),'data':data,'runs':runs,'aggregate':aggregate}
    Path(__file__).with_name('results.json').write_text(json.dumps(result,indent=2)+'\n')
    print(json.dumps({k:v for k,v in result.items() if k not in ['data','runs']},indent=2))
    print('teacher_accuracy', [r['teacher']['accuracy'] for r in runs])
    print('PASS: frozen teacher, shared student initialization, gradient and KL checks')


if __name__=='__main__':
    main()