"""CPU necessary-condition tests; no remote training or optimizer updates."""

import ast
import copy
import hashlib
import json
import subprocess
import sys
import time
from pathlib import Path
from types import SimpleNamespace

import torch
import torch.nn.functional as F

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, '/home/ubuntu/repos/research-exchange')
from research_exchange.correctness import parity_errors, validate


def digest(path):
    return hashlib.sha256(path.read_bytes()).hexdigest()


def extract(path, names, namespace):
    tree = ast.parse(path.read_text())
    nodes = [n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name in names]
    assert {n.name for n in nodes} == set(names)
    for node in nodes:
        node.decorator_list = []
    source = 'from __future__ import annotations\n' + '\n\n'.join(ast.unparse(n) for n in nodes) + '\n'
    exec(compile(source, str(path), 'exec'), namespace)  # noqa: S102 -- only pinned reference function ASTs
    return source


def flatten(x):
    return x.detach().reshape(-1).tolist()


def errors(reference, actual):
    return parity_errors(flatten(reference), flatten(actual))


def configured_is_loss(logp, old, adv, mask):
    return -(torch.exp(logp - old) * adv * mask).sum() / mask.sum()


def logprobs(weights, hidden, targets):
    return F.log_softmax(hidden @ weights, dim=-1).gather(-1, targets[..., None]).squeeze(-1)


def main():
    started = time.perf_counter()
    torch.set_num_threads(1)
    torch.manual_seed(0)
    ns = {'torch': torch, 'F': F, 'RL_LOSSES': {'importance_sampling'}}
    sdpo_path = ROOT / 'SDPO/verl/trainer/ppo/core_algos.py'
    masked_source = extract(ROOT / 'SDPO/verl/utils/torch_functional.py', ['masked_sum'], ns)
    ns['verl_F'] = SimpleNamespace(masked_sum=ns['masked_sum'])
    sdpo_source = extract(sdpo_path, ['agg_loss', 'compute_self_distillation_loss'], ns)
    stock_source = extract(ROOT / 'spindle/src/spindle/backends/megatron_runtime/common/forward_backward.py', ['_loss'], ns)
    adv_source = extract(ROOT / 'tinker-cookbook/tinker_cookbook/rl/data_processing.py', ['compute_advantages'], ns)
    (ROOT / 'reference-extracts.py').write_text(masked_source + sdpo_source + stock_source + adv_source)

    hidden = torch.randn(4, 7, 5, dtype=torch.float32)
    weights = torch.randn(5, 11, dtype=torch.float32) * 0.2
    targets = torch.randint(0, 11, (4, 7))
    mask = torch.tensor([[0, 1, 1, 1, 1, 0, 0], [0, 1, 1, 0, 0, 0, 0],
                         [0, 1, 1, 1, 1, 1, 0], [0, 1, 1, 1, 0, 0, 0]], dtype=torch.float32)
    old = logprobs(weights, hidden, targets).detach()
    delta = 0.3 * torch.sin(torch.arange(28, dtype=torch.float32).reshape(4, 7) + 0.4)
    teacher = (old + delta).detach()
    rewards = torch.tensor([0.2, 0.4, 0.7, 0.9], dtype=torch.float32)
    fixture = {
        'seed': 0, 'dtype': 'float32', 'hidden': hidden.tolist(), 'weights': weights.tolist(),
        'targets': targets.tolist(), 'mask': mask.tolist(), 'old_logprobs': old.tolist(),
        'teacher_logprobs': teacher.tolist(), 'rewards': rewards.tolist(),
        'optimizer_state': None, 'optimizer_explanation': 'Pure loss/gradient test; no optimizer state or updates.',
        'description': 'Synthetic linear-softmax policy, full 5x11 parameter gradient, variable-length masks; not real rollout or activation evidence.',
    }
    fixture_path = ROOT / 'fixed-fixture.json'
    fixture_path.write_text(json.dumps(fixture, indent=2) + '\n')

    group = SimpleNamespace(get_total_rewards=lambda: rewards.tolist())
    reference_adv = ns['compute_advantages']([group])[0][:, None].expand_as(old).detach()
    custom_adv = (rewards - rewards.mean())[:, None].expand_as(old).detach()
    cases = []
    gradients = []
    for name, stale, packed in [('on_policy_unpacked', False, False), ('on_policy_packed', False, True),
                                ('stale_unpacked', True, False), ('stale_packed', True, True)]:
        w = weights.clone()
        if stale:
            w += 0.08 * torch.cos(torch.arange(w.numel()).reshape_as(w))
        w.requires_grad_()
        lp = logprobs(w, hidden, targets)
        m, o, t, ca, ra = mask, old, teacher, custom_adv, reference_adv
        if packed:
            lp, m, o, t, ca, ra = [x.reshape(1, -1) for x in (lp, m, o, t, ca, ra)]

        def measure(loss, w=w):
            return torch.autograd.grad(loss, w, retain_graph=True)[0]

        new_grpo = configured_is_loss(lp, o, ca, m)
        stock_sum, stock_count = ns['_loss'](lp.reshape(-1), {'loss_name': 'importance_sampling', 'sampling_logprobs': o,
                                                  'advantages': ra, 'loss_mask': m})
        stock_grpo = stock_sum / stock_count
        new_sdpo = configured_is_loss(lp, o, (t - o).detach(), m)
        cfg = SimpleNamespace(full_logit_distillation=False, alpha=1.0, is_clip=None)
        ref_sdpo, _ = ns['compute_self_distillation_loss'](lp, t, m, cfg, loss_agg_mode='token-mean')
        cfg_is = SimpleNamespace(full_logit_distillation=False, alpha=1.0, is_clip=2.0)
        ref_sdpo_is, _ = ns['compute_self_distillation_loss'](lp, t, m, cfg_is, old_log_probs=o, loss_agg_mode='token-mean')
        zero_new = configured_is_loss(lp, o, torch.zeros_like(o), m)
        zero_ref, _ = ns['compute_self_distillation_loss'](lp, lp.detach(), m, cfg, loss_agg_mode='token-mean')
        ng, rg = measure(new_sdpo), measure(ref_sdpo)
        case = {
            'name': name, 'loss_tokens': int(m.sum()),
            'grpo_loss': errors(stock_grpo, new_grpo), 'grpo_gradient': errors(measure(stock_grpo), measure(new_grpo)),
            'sdpo_zero_loss': errors(zero_ref, zero_new), 'sdpo_zero_gradient': errors(measure(zero_ref), measure(zero_new)),
            'sdpo_loss': errors(ref_sdpo, new_sdpo), 'sdpo_gradient': errors(rg, ng),
            'sdpo_is_clip_2_loss': errors(ref_sdpo_is, new_sdpo),
            'sdpo_is_clip_2_gradient': errors(measure(ref_sdpo_is), ng),
            'values': {'configured_is_loss': new_sdpo.item(), 'authors_sampled_rkl_loss': ref_sdpo.item(),
                       'authors_sampled_rkl_is_clip_2_loss': ref_sdpo_is.item()},
        }
        cases.append(case)
        gradients.append({'case': name, 'configured': flatten(ng), 'reference': flatten(rg),
                          'reference_is_clip_2': flatten(measure(ref_sdpo_is))})

    spec = json.loads((ROOT / 'spec.json').read_text())
    effective = copy.deepcopy(spec)
    cookbook_sha = subprocess.check_output(['git', '-C', str(ROOT / 'tinker-cookbook'), 'rev-parse', 'HEAD'], text=True).strip()
    sdpo_sha = subprocess.check_output(['git', '-C', str(ROOT / 'SDPO'), 'rev-parse', 'HEAD'], text=True).strip()
    for plan in effective['correctness_preflight']['paths']:
        plan['baseline_revision'] = cookbook_sha
        plan['reference_revision'] = sdpo_sha if plan['id'].startswith('sdpo') else cookbook_sha

    first = cases[0]
    reports = []
    for plan in effective['correctness_preflight']['paths']:
        sdpo = plan['id'].startswith('sdpo')
        reports.append({
            'id': plan['id'], 'fixed_input': {'sha256': digest(fixture_path), 'description': 'fixed-fixture.json: CPU float32 synthetic linear-softmax policy, identical weights, token identities, masks; full parameter gradients. No optimizer or remote execution.'},
            'baseline_revision': plan['baseline_revision'], 'reference_revision': plan['reference_revision'],
            'implementation_revision': digest(Path(__file__)), 'precision': 'CPU torch.float32; no mixed precision; not a Spindle backend parity test',
            'normalization': 'sum of masked token losses / 14 loss tokens; no reward mixing; no clipping in specified IS path',
            'errors': { 'baseline_loss': first['sdpo_zero_loss' if sdpo else 'grpo_loss'],
                        'baseline_gradient': first['sdpo_zero_gradient' if sdpo else 'grpo_gradient'],
                        'reference_loss': first['sdpo_loss' if sdpo else 'grpo_loss'],
                        'reference_gradient': first['sdpo_gradient' if sdpo else 'grpo_gradient']},
            'lr_scaling': [],
            'scope': 'Necessary-condition offline test only; cannot certify server precision, init, or optimizer.',
        })
    report = {'required': True, 'paths': reports}
    effective_validation = validate(effective, report)
    original_validation = validate(spec, report)
    output = {
        'cases': cases, 'effective_validation': effective_validation, 'immutable_spec_validation': original_validation,
        'torch_version': torch.__version__, 'seconds': time.perf_counter() - started,
        'optimizer_updates': 0, 'training_clients_created': 0,
        'source_hashes': {'sdpo_core_algos': digest(sdpo_path), 'reference_extracts': digest(ROOT / 'reference-extracts.py')},
        'reference_config': {'full_logit_distillation': False, 'alpha': 1.0, 'is_clip': None},
        'additional_reference_config': {'full_logit_distillation': False, 'alpha': 1.0, 'is_clip': 2.0},
    }
    for name, value in [('effective-spec-pinned.json', effective), ('correctness-preflight.json', report),
                        ('correctness-results.json', output), ('full-gradients.json', gradients)]:
        (ROOT / name).write_text(json.dumps(value, indent=2, allow_nan=False) + '\n')
    print(json.dumps(output, indent=2, allow_nan=False))


if __name__ == '__main__':
    main()