import json, math, random
import numpy as np
import torch

SEED = 137
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)

def pce_components(teacher, student, sigma=0.35, eps=1e-8):
    # teacher/student: [n,2] birth-death coordinates
    dist2 = ((teacher[:, None, :] - student[None, :, :]) ** 2).sum(-1)
    sim = torch.exp(-dist2 / (2 * sigma * sigma))
    a = sim.max(dim=1).values
    persistence = (teacher[:, 1] - teacher[:, 0]).clamp_min(0)
    p = persistence / (persistence.sum() + eps)
    explained = p * a
    unexplained = 1.0 - explained.sum()
    loss = -(p * torch.log(p * a + eps)).sum()
    entropy = -(p * torch.log(p + eps)).sum()
    return loss, entropy, unexplained, a, p

def chamfer_loss(teacher, student):
    d = torch.cdist(teacher, student).pow(2)
    return d.min(1).values.mean() + d.min(0).values.mean()

def run(method, teacher, init, steps=500, lr=0.04):
    student = torch.nn.Parameter(init.clone())
    opt = torch.optim.Adam([student], lr=lr)
    history=[]
    for step in range(steps):
        opt.zero_grad()
        if method == 'pce':
            loss, _, _, _, _ = pce_components(teacher, student)
            # The entropy term is constant, but retaining it matches the proposed excess objective.
            objective = loss
        else:
            objective = chamfer_loss(teacher, student)
        objective.backward(); opt.step()
        if step in (0, 99, 249, 499):
            with torch.no_grad():
                pc = pce_components(teacher, student)
                history.append((step + 1, float(pc[2]), float(pc[0]), float(objective)))
    with torch.no_grad():
        loss, ent, unexplained, a, p = pce_components(teacher, student)
    return student.detach(), history, (loss.item(), ent.item(), unexplained.item(), a.numpy(), p.numpy())

def main():
    # Four teacher events with unequal importance and only two student events.
    teacher = torch.tensor([[0.10,0.90], [0.18,0.52], [0.62,0.96], [0.38,0.48]], dtype=torch.float32)
    init = torch.tensor([[0.25,0.25], [0.75,0.75]], dtype=torch.float32)

    # Core numerical verification: valid measure and excess identity for random diagrams.
    torch.manual_seed(11)
    T = torch.rand(7,2); T[:,1] = T[:,0] + torch.rand(7)
    S = torch.rand(3,2); S[:,1] = S[:,0] + torch.rand(3)
    l,e,u,a,p = pce_components(T,S,sigma=.27)
    identity_error = abs((l-e).item()) # overwritten below; calculate stated excess exactly
    excess_error = abs((l-e).item() - (-(p*torch.log(a+1e-8)).sum()).item())
    validity = float((p >= 0).all() and (a >= 0).all() and (a <= 1+1e-6).all() and (u >= -1e-6).item() and (u <= 1+1e-6).item())

    pce_student, pce_hist, pce_final = run('pce', teacher, init)
    cham_student, cham_hist, cham_final = run('chamfer', teacher, init)
    # Repeated starts test whether the directional signal survives initialization.
    repeat = []
    for seed in (3, 19, 71):
        g = torch.Generator().manual_seed(seed)
        start = torch.rand((2, 2), generator=g) * 0.8 + 0.1
        ps = run('pce', teacher, start)[2]
        cs = run('chamfer', teacher, start)[2]
        repeat.append({'seed': seed, 'pce_unexplained': ps[2],
                       'chamfer_unexplained': cs[2],
                       'pce_excess': ps[0]-ps[1],
                       'chamfer_excess': cs[0]-cs[1]})
    out = {
      'seed': SEED, 'teacher_points': teacher.tolist(), 'student_count': 2,
      'repeated_starts': repeat,
      'core_check': {'probability_valid': bool(validity), 'excess_identity_abs_error': excess_error,
                     'unexplained_mass': u.item(), 'p_sum': p.sum().item(), 'explained_sum': (p*a).sum().item()},
      'pce': {'student_points': pce_student.tolist(), 'trajectory': pce_hist,
              'final_loss': pce_final[0], 'teacher_entropy': pce_final[1],
              'final_excess': pce_final[0]-pce_final[1], 'unexplained_mass': pce_final[2],
              'explanation_scores': pce_final[3].tolist()},
      'chamfer_control': {'student_points': cham_student.tolist(), 'trajectory': cham_hist,
              'pce_loss_evaluated': cham_final[0], 'teacher_entropy': cham_final[1],
              'final_excess': cham_final[0]-cham_final[1], 'unexplained_mass': cham_final[2],
              'explanation_scores': cham_final[3].tolist()},
    }
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))

if __name__ == '__main__': main()