import dataclasses
import hashlib
import json
from pathlib import Path
import subprocess

import numpy as np
import torch
from research_exchange.correctness import parity_errors
from gs1_reference import reference_baselines
from rx_gs1 import ROOT, advantages, dataset, datums, kernel_baselines, sample_batch, service, write_json
from tinker_cookbook.rl.data_processing import assemble_training_data, compute_advantages


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


def objective(data, outputs):
    losses, gradients = [], []
    for d, o in zip(data, outputs, strict=True):
        new = np.array(o["logprobs"].data, dtype=np.float64)
        old = np.array(d.loss_fn_inputs["logprobs"].data, dtype=np.float64)
        adv = np.array(d.loss_fn_inputs["advantages"].data, dtype=np.float64)
        grad = -np.exp(new - old) * adv
        losses.append(grad.sum())
        gradients.extend(grad.tolist())
    return [float(np.sum(losses))], gradients


async def preflight(tc=None):
    text = (ROOT / "launch-instructions.md").read_text()
    spec, _ = json.JSONDecoder().raw_decode(text[text.index("{", text.index("Experiment spec:")):])
    write_json(ROOT / "spec.json", spec)
    cookbook_sha = subprocess.check_output(["git", "-C", str(ROOT / "cookbook"), "rev-parse", "HEAD"], text=True).strip()
    reference_commit = subprocess.check_output(["git", "-C", str(ROOT), "log", "-1", "--format=%H", "--", "gs1_reference.py"], text=True).strip()
    reference_revision = f"git:{reference_commit};sha256:{sha(ROOT / 'gs1_reference.py')}"
    effective = json.loads(json.dumps(spec))
    path = effective["correctness_preflight"]["paths"][0]
    path.update(baseline_revision=cookbook_sha, reference_revision=reference_revision)
    effective["correctness_preflight"]["status"] = "pinned for measured preflight; see preflight.json"
    write_json(ROOT / "effective-spec.json", effective)
    ds, _ = await dataset(8, 8)
    sc = service() if tc is None else None
    status = "errored"
    try:
        if tc is None:
            tc = await sc.create_lora_training_client_async(base_model="Qwen/Qwen3.5-2B", rank=32)
            saved = await tc.save_state_async("rx-gs1-init-run_5207121d26")
            response = await saved.result_async()
            write_json(ROOT / "init-save-response.json", dataclasses.asdict(response))
        sampler = await tc.save_weights_and_get_sampling_client_async()
        groups, features = await sample_batch(sampler, ds.get_batch(0))
        stock = compute_advantages(groups)
        disabled, _, _ = advantages(groups, "grpo_gs8", features, [])
        rewards = np.array([r for g in groups for r in g.get_total_rewards()], dtype=np.float64)
        b, neighbors = kernel_baselines(features, rewards, len(rewards))
        ref_b, _ = reference_baselines(features, rewards, len(rewards))
        # The enabled loop uses singleton groups. Preserve the identical sampled trajectories.
        singleton_groups = [dataclasses.replace(g, trajectories_G=[g.trajectories_G[i]], final_rewards_G=[g.final_rewards_G[i]], metrics_G=[g.metrics_G[i]]) for g in groups for i in range(len(g.trajectories_G))]
        enabled, actual_b, _ = advantages(singleton_groups, "gs1_kernel", features, [])
        reference_adv = [torch.tensor([r - base], dtype=torch.float64) for r, base in zip(rewards, ref_b)]
        stock_data, _ = datums(groups, stock)
        disabled_data, _ = datums(groups, disabled)
        enabled_data, _ = datums(singleton_groups, enabled)
        reference_data, _ = datums(singleton_groups, reference_adv)
        masked_data, _ = assemble_training_data(groups, stock, include_provenance=False)
        fixture = dict(features=features, rewards=rewards.tolist(), group_sizes=[8] * 8,
            stock_datums=stock_data, enabled_datums=enabled_data,
            disabled_datums=disabled_data, reference_datums=reference_data,
            stock_masks=[d.loss_fn_inputs["mask"].data for d in masked_data],
            baselines=actual_b.tolist(), reference_baselines=ref_b.tolist())
        write_json(ROOT / "fixed-batch.json", fixture)
        print("Fixed batch saved; one forward, no backward or optimizer calls.", flush=True)
        forward_future = await tc.forward_async(stock_data, loss_fn="importance_sampling")
        output = await forward_future.result_async()
        write_json(ROOT / "trainer-forward.json", dict(metrics=output.metrics,
            outputs=output.loss_fn_outputs))
        baseline_loss, baseline_grad = objective(stock_data, output.loss_fn_outputs)
        actual_loss, actual_grad = objective(disabled_data, output.loss_fn_outputs)
        ref_loss, ref_grad = objective(reference_data, output.loss_fn_outputs)
        gs_loss, gs_grad = objective(enabled_data, output.loss_fn_outputs)
        errors = dict(baseline_loss=parity_errors(baseline_loss, actual_loss),
            baseline_gradient=parity_errors(baseline_grad, actual_grad),
            reference_loss=parity_errors(ref_loss, gs_loss), reference_gradient=parity_errors(ref_grad, gs_grad))
        report = dict(required=True, paths=[dict(id=path["id"], baseline_revision=cookbook_sha,
            reference_revision=reference_revision, implementation_revision=f"sha256:{sha(ROOT / 'rx_gs1.py')}",
            fixed_input=dict(sha256=sha(ROOT / "fixed-batch.json"), description="fixed-batch.json: seed-0 first 8 GSM8K train prompts x8, sampled init rank32; trainer-forward.json: one forward on never-updated init client; no optimizer state or updates"),
            precision="numpy float64 analytic loss/full dL/dlogprobs, trainer/sampler returned logprobs, stock advantages float32",
            normalization="negative token sum of exp(trainer_logprob-sampler_logprob)*advantage, no token/batch divisor; prompt token advantages zero; no packing; full per-token logprob gradient",
            errors=errors)])
        write_json(ROOT / "preflight.json", report)
        write_json(ROOT / "measurements.json", dict(errors=errors, kernel_baseline_max_abs=float(np.max(np.abs(actual_b-ref_b))),
            reward_values=np.unique(rewards).tolist(), nonzero_stock_grad=int(np.count_nonzero(baseline_grad)),
            nonzero_gs1_grad=int(np.count_nonzero(gs_grad)), server_metrics=output.metrics, local_stock_loss=baseline_loss))
        print(json.dumps(errors, indent=2), flush=True)
        status = "success"
    finally:
        if sc is not None:
            await sc.close(status)
            print("Service session closed; clients released.", flush=True)


if __name__ == "__main__":
    import asyncio
    from rx_gs1 import guard_grader
    guard_grader()
    asyncio.run(preflight())